Создание агента-ассистента с знаниями на Databricks с использованием LangGraph

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

Пошаговая инструкция включает набор нотбуков, сгруппированных по папкам: основы, конвейер RAG и модуль памяти.

Подготовка рабочей среды

Полноценный агент включает векторные индексы, инструменты и журналы логов, и им сначала необходимо создать организованную структуру хранения. В Databricks слоем управления является Unity Catalog (UC), который контролирует все ресурсы в рамках трехуровневой структуры имен. Последовательное выполнение ноутбуков в папке 00_foundations создает все необходимые ресурсы, от которых зависят последующие шаги, включая смесь структурированных и неструктурированных данных, к которым будет подключаться агент.

Память будет храниться в Lakebase — автомасштабируемой базе данных, совместимой с Postgres. Чтобы позже можно было без проблем создавать таблицы, ноутбук настройки также создает проект Lakebase и ветку для их хранения.

Первый минимальный агент

Далее идут ноутбуки в папке 01_rag_pipeline, начиная с 0_your_first_agent_on_databricks.

Каждый агент строится на основе базовой модели. В Databricks Model Serving перечислены доступные модели, причем этот список может отличаться в зависимости от настроек вашей рабочей среды. Конечные точки совместимы с форматом API OpenAI, но для тесной интеграции с LangGraph лучше использовать официальную библиотеку databricks-langchain.

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

Вызывается агент, и появляется трейс, в котором указаны входной вопрос, результат работы модели и каждый логический шаг между ними. Это функция MLflow, который нативно интегрируется с langgraph. Один вызов mlflow.langchain.autolog() в начале ноутбука включает эти логи и диагностическую информацию. Более подробно о MLflow рассказано позже.

В этой первой версии вопрос пользователя просто передаётся модели. Настоящий агент добавляет системное промпт, которое задаёт тон, формирует результат и предоставляет контекст, например, для чего предназначен ассистент и на что ему следует отказаться отвечать. Поэтому во второй версии в граф вставляется системное промпт.

Добавление механизма поиска

Одним из наиболее распространенных типов агентов является агент по поиску знаний, также называемый агентом RAG. Он помогает людям искать среди большого количества документов и отвечать на вопросы о них. Чтобы предоставить агенту такую способность, ему необходим индекс векторного поиска — по сути, база данных для неструктурированного текста.

В рабочей тетради 01_building_a_doc_store_on_vector_search рассматривается процесс создания такого индекса. Прежде чем индекс можно будет привязать к агенту, он должен существовать. В этой части описываются основные шаги, при этом опускается настройка и оптимизация.

Для создания индекса требуется выполнить три действия:

• Анализ означает открытие документа и чтение его содержимого.
• Разбиение на фрагменты подразумевает использование инструмента для разделения документа на части.
• Генерация векторов означает преобразование каждого фрагмента в семантический вектор, хранящийся в индексе.

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

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

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

На первых двух этапах в ноутбуке файлы формата Markdown, сгенерированные во время настройки, загружаются и делятся на перекрывающиеся блоки по 800 символов каждый, после чего они загружаются в Databricks Vector Search. Библиотека databricks_langchain включает готовый инструмент VectorSearchRetrieverTool, который соединяет индекс с агентом langgraph.

После того как индекс создан и протестирован, на третьем этапе происходит переопределение агента. Он выглядит почти так же, как предыдущие версии, но с одним дополнением: элементом ToolNode, который принимает список инструментов. В данном случае этот список содержит только VectorSearchRetrieverTool.

В получившейся схеме вопрос поступает и проходит через шаблон запроса. Модель видит запрос вместе с инструментами и их описаниями, после чего решает, следует ли обращаться к индексу. Она может делать это несколько раз и отвечает пользователю только тогда, когда находит искомое.

Это рабочий агент типа RAG. Граф остается компактным, при этом единственным дополнением является узел инструмента, что показывает, насколько мало структуры требуется для цикла поиска. Однако у помощника для работы с знаниями нужен ещё один элемент — память.

Предоставление памяти агенту

На текущем этапе агент совершенно не обладает памятью. Попытка задать дополнительный вопрос, отсылающий к предыдущему ответу, просто терпит неудачу, поскольку модель не помнит, что говорилось мгновение назад.

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

Существует два вида памяти. Краткосрочная память охватывает только текущую активную сессию. Долгосрочная память сохраняет информацию между сессиями. Почти каждому приложению необходима первая форма памяти, а нужна ли вторая — зависит от конкретного случая использования.

Краткосрочная память

Папка 02_memory открывается вместе с 01_short_term_memory, где добавляется контрольный указатель. На шагах 1–3 создается StateGraph, идентичный предыдущему, а на шаге 4 показано, насколько неуклюжей становится беседа без памяти.

Создание контрольного указателя, связанного с Lakebase, восстанавливает преемственность, а databricks_langchain предоставляет встроенный класс CheckpointSaver для его создания и привязки к графу.

Агент — это сервис, используемый многими пользователями, поэтому сохранённые беседы должны иметь владельца. Для этой цели используется thread_id. В производственном приложении фронтенд должен генерировать уникальный thread_id для каждой активной сессии, чтобы каждый пользователь видел только свою историю чатов.

К шагу 6 контрольный указатель уже установлен, и беседа проходит естественно.

Долгосрочная память

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

Для краткосрочной памяти достаточен CheckpointSaver. Для долгосрочной памяти требуется определённый дизайн. Обычно не нужно записывать каждый разговор на будущее. Стоит сохранять такие факты, как имя или место рождения, а также предпочтения пользователя.

В записной книжке 02_long_term_memory описываются технические детали. Шаги с 1 по 4 предусматривают создание агента RAG с краткосрочной памятью, а шаг 5 демонстрирует, что предпочтения исчезают при изменении сессии и thread_id.

На шаге 6 вводится объект DatabricksStore, который сохраняет данные в Lakebase. Эта часть проста; настоящая задача проектирования заключается в определении того, что следует сохранять. На шаге 7 используется простой подход: в хранилище записывается небольшое количество жестко заданных фактов.

Чтобы использовать их, на шаге 8 создается заново агент. Ранее узел агента представлял собой простую функцию, вызывающую метод invoke у основной модели. Теперь это более сложная функция: она принимает ID пользователя (жестко заданный в демонстрации), загружает факты этого пользователя из DatabricksStore и объединяет их в контекст, отправляемый модели. Очевидно, что такой подход требует тщательного проектирования для стабильной работы.

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

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

Как дальше

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