Створення агента знанневого асистента у Databricks за допомогою LangGraph

Базовий користувацький агент у Databricks може розвиватися від простого виклику LLM до інструменту знань, який шукає інформацію в документах та пам’ятає користувачів; кожен крок у цьому процесі є незначним. Готовий агент створюється за допомогою фреймворку 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 організовує це шляхом моделювання агента у вигляді графа. Вузли містять функції та виклики LLM, а ребра визначають, як контекст переміщується між ними. Вихідний граф є дуже простим.

Викликаєтьте агента, і з’являється запис, у якому вказано запит користувача, результат роботи моделі та кожен логічний крок між ними. Це забезпечується 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 у базовій моделі; тепер це більш складна функція. Вона отримує ідентифікатор користувача (заданий як константа для демонстрації), завантажує факти цього користувача з DatabricksStore та об’єднує їх у контекст, який надсилається до моделі. Ця схема однозначно вимагає ретельного проектування, щоб бути ефективною.

Більш гнучкий підхід дозволяє самій моделі вирішувати, коли зберігати, видаляти та отримувати інформацію з пам’яті. Крок 12 реалізує базову версію: три нові інструменти взаємодіють з модулем пам’яті, а системний запит змінюється, щоб повідомити модель про них та заохотити до їх використання.

Результатом є повноцінний асистент, причому кожну з його функцій можна перевірити окремо, оскільки вона була реалізована через окремий ноутбук та окрему зміну в структурі графа. Він шукає документи за допомогою векторного індексу, відстежує поточну розмову за допомогою короткострокової пам’яті та запам’ятовує деталі про користувачів між сеансами за допомогою довгострокової пам’яті.

Яким буде наступний крок

Шлях розвитку пройшов від бота, який передавав запитання безпосередньо до LLM, до системи з шаблоном запиту, короткостроковою та довгостроковою пам’яттю, а також функцією векторного пошуку для отримання контексту. Наступним кроком є впровадження цієї системи у промислове використання: оцінка якості агента, додавання нових функцій та розгортання її як сервісу для кінцевих користувачів.