Привет, Хабр! Меня зовут Владимир и это вторая часть материала о трассировке LLM-агентов. В первой части мы настроили инфраструктуру: подняли LangFuse, организовали трассировку и научились управлять промптами как кодом. Если вы ещё не читали — рекомендую начать с неё.

В этой части перейдём от теории к практике: соберём агента, который пишет сказки. В графе будут задействованы инструменты, условные переходы и циклы обратной связи.

Поехали. Граф будет состоять из четырёх узлов - Архитектор, Писатель, Критик, Редактор и трёх инструментов для генерации тематики сказки. Ориентировочная схема графа

План такой - Узел Архитектор с помощью инструментов накидывает тему сказки для писателя, узел Писатель пишет черновик, дальше узел Критик критикует и возвращает Писателю. Писатель исправляет на основании замечаний, опять передаёт Критику и так, пока Критику не понравится (ну или количество попыток закончится). В конце узел Редактор полирует текст.

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

import random
from langchain_core.tools import tool
@tool
def generate_hero() -> str:
    """Генерирует случайного героя сказки из шаблонов. Используй, когда нужно создать главного персонажа."""
    templates = [
        "молодой кузнец с золотыми руками",
        "старая ведьма с добрым сердцем",
        "говорящий волк, изгнанный из стаи",
        "принцесса, сбежавшая из замка",
        "слепой музыкант с волшебной флейтой",
        "вор-неудачник не умеющий подкрадываться",
        "рыцарь, закованный в ржавые доспехи",
        "девочка, умеющая разговаривать с тенями"
    ]
    return random.choice(templates)

Как-бы всё. Дальше LangChain сам сделает имя инструмента из имени функции и описание из докстринга. В этом способе важно описать в докстринге суть инструмента, т.к. на основании этого описания LLM будет принимать решение о необходимости вызова данного инструмента. Если есть параметры - надо обязательно аннотировать, опять же для правильного вызова инструмента с помощью LLM (в том числе возвращаемые значения).

Способ посложнее:

from langchain.tools import BaseTool

class GenerateHeroTool(BaseTool):
    name: str = 'generate_hero'
    description: str = 'Генерирует случайного героя сказки из шаблонов. Используй, когда нужно создать главного персонажа.'

    def __init__(self):
        super().__init__()
        self._templates = [...]

    def _run(self) -> str:
        return random.choice(self._templates)

    async def _arun(self) -> str:
        return self._run()

В данном варианте мы наследуемся от langchain.tools.BaseTool, вручную прописываем имя инструмента и его описание. Также надо реализовать синхронный и асинхронный методы запуска инструмента. Данный способ даёт больше контроля над инструментом, но гораздо сложнее. Пользоваться им сейчас мы, конечно же, не будем.

Реализацию инструментов generate_villain, generate_location приводить не буду, т.к. это тот-же generate_hero, но с другими фразами. После описания наших инструментов необходимо собрать их в одну кучку, для дальнейшего построения узла:

tools_list = [generate_hero, generate_villain, generate_location]

Теперь займёмся Состоянием графа. Оно, кроме истории сообщений, должно содержать поля для хранения запроса пользователя, задания от Архитектора к Писателю, черновика Писателя, замечаний Критика. Ну и счётчик итераций, что-бы, если что, прервать Критика. Создадим в файле storyteller\state.py наш StoryState, используя ранее описанные фишечки:

class StoryState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    user_topic: str
    task_for_writer: str | None
    draft: str | None
    critique: str | None
    iterations: Annotated[int, operator.add]

Перейдём к созданию Графа. В файле storyteller\graph.py начнём описывать наш класс-граф:

from .tools import tools_list
class StoryTeller:
    def __init__(self, llm):
        self.llm = llm
        self.llm_with_tools = self.llm.bind_tools(tools_list)
        self.graph = self.compile_graph()

В конструктор будем передавать наш экземпляр ChatOpenAI. Запоминаем его для использования в узлах. Дополнительно сразу сконфигурируем нашу LLM для работы с инструментами (для Архитектора). Метод self.compile_graph() рассмотрим после узлов. Собственно, к узлам.

Начнём с входного. Его смысл - просто положить в историю сообщений запрос пользователя:

from langchain_core.messages import HumanMessage
from .state import StoryState

class StoryTeller:
    # остальные поля из предыдущего листинга
    @staticmethod
    def start_node(state: StoryState) -> dict[str, Any]:
        return {'messages': [HumanMessage(content=f"Тема сказки: {state['user_topic']}")]}

Следующий узел - Архитектор. В нём, как и в последующих узлах, основная работа приходится на промпты, хранящиеся в Langfuse. Сначала промпт, который управляет Архитектором. Его делаем чатом с двумя сообщениями - Системный промпт и User промпт (с переменной для подстановки входной задачи):

{"role": "system", "content": "Ты — Архитектор Сказок. Собираешь элементы для сказки.\n\nУ тебя есть Инструменты:\n- generate_hero() — предоставляет героя для сказки. Вызывай БЕЗ параметров\n- generate_villain() — предоставляет злодея для сказки. Вызывай БЕЗ параметров\n- generate_location() — предоставляет локацию для сказки. Вызывай БЕЗ параметров\n\nИНСТРУКЦИЯ:\n1. Вызывай инструменты для получения сущности\n2. Изучай результаты и решай — достаточно элементов или вызвать ещё\n3. Когда готов — напиши ЗАДАНИЕ ДЛЯ ПИСАТЕЛЯ. Опиши Сюжет, Места, Персонажей, Стиль сказки"}
{"role": "user", "content": "Тема пользователя: {{user_topic}}"}

В коде узла получаем из Langfuse промпт, при компиляции подставляем в него наш запрос state['user_topic']

def architect_node(self, state: StoryState):
    prompt = settings.langfuse.client.get_prompt(name='architect_prompt').compile(user_topic=state['user_topic'])
    response = self.llm_with_tools.invoke(prompt)
    if hasattr(response, 'tool_calls') and response.tool_calls:
        return {'messages': [response]}
    return {'messages': [response], 'task_for_writer': response.content}

У узла должно быть два возможных выхода - на инструменты и на Писателя. Это реализует условие if. Если ответ LLM содержит tool_calls (запрос модели на вызов инструментов), то просто сохраняем в историю ответ модели. Если нет, то это итоговое задание для писателя и мы дополнительно сохраняем его в поле Состояния графа. Далее узлы Писателя, Критика и Редактора. В них тоже промпт-инжинеринг в чистом виде (через Langfuse):

{"role": "system", "content": "Ты — Сказочник. Напиши сказку по заданию.\n\nЗАДАНИЕ:\n{{task_for_writer}}\n\nТЕКУЩАЯ ВЕРСИЯ:\n{{draft}}\n\nЗАМЕЧАНИЯ КРИТИКА:\n{{critique}}\n\nИНСТРУКЦИЯ:\n1. Пиши живым сказочным языком, с диалогами и описаниями\n2. Если есть замечания критика — учти их\n3. Объём: 300-500 слов\n\nНапиши полную сказку:"}

{"role": "system", "content": "Ты — Литературный Критик. Оцени сказку.\n\nТЕКУЩАЯ ВЕРСИЯ:\n{{draft}}\n\nЗАМЕЧАНИЯ КРИТИКА:\n{{critique}}\n\nИНСТРУКЦИЯ:\n1. Оцени: конфликт, структуру, язык\n2. Если всё хорошо — вердикт \"ГОТОВО\"\n3. Если есть проблемы — вердикт \"ДОРАБОТАТЬ\" + конкретные замечания\n\nФОРМАТ:\nВЕРДИКТ: <ГОТОВО или ДОРАБОТАТЬ>\nКОММЕНТАРИИ: <замечания>"}

{"role": "system", "content": "Ты — Главный Редактор. Финальная полировка.\n\nЧЕРНОВИК:\n{{draft}}\n\nПРОАНАЛИЗИРУЙ текст сказки, отшлифуй его и верни ФИНАЛЬНУЮ версию"}

Соответственно, код также отличается только подставляемыми переменными в промпт. Для избежания зацикливания вызовов Писатель-Критик в Состоянии есть поле для подсчета итераций. Значение счетчика обновляет Критик:

def writer_node(self, state: StoryState) -> dict[str, Any]:
    prompt = settings.langfuse.client.get_prompt(name='writer_prompt').compile(
        task_for_writer=state['task_for_writer'],
        draft=state['draft'] if state.get('draft') else 'Отсутствует',
        сritique=state['critique'] if state.get('critique') else 'Отсутствует'
    )
    response = self.llm.invoke(prompt)
    return {'draft': response.content, 'messages': [response]}


def critic_node(self, state: StoryState) -> dict[str, Any]:
    prompt = settings.langfuse.client.get_prompt(name='critic_prompt').compile(
        draft=state['draft'],
        critique=state['critique'] if state.get('critique') else 'Отсутствует'
    )
    response = self.llm.invoke(prompt)
    return {'critique': response.content, 'messages': [response], 'iterations': 1}


def corrector_node(self, state: StoryState) -> dict[str, Any]:
    prompt = settings.langfuse.client.get_prompt(name='corrector_prompt').compile(draft=state['draft'])
    response = self.llm.invoke(prompt)
    return {'messages': [response]}

Теперь метод маршрутизации. Он будет определять, куда идти из узла Критик - назад к Писателю или к Редактору. Так как в промпте у Критика задали ключевое слово ГОТОВО, то его и будем использовать как признак передачи текста Редактору. Чтобы избежать бесконечных правок, так же будем контролировать, что счётчик меньше порогового MAX_ITERS:

@staticmethod
def route_after_critic(state: StoryState) -> Literal['writer', 'corrector']:
    if state['iterations'] > MAX_ITERS or 'ГОТОВО' in state['critique'].upper():
        return 'corrector'
    return 'writer'

Всё. База готова, переходим к созданию графа. Сначала создадим специальный инструментный узел - ToolNode. Он предназначен для выполнения наших инструментов по запросу модели и возврата результата выполнения. В узел надо передать список используемых инструментов. Если привести простую аналогию, то наша self.llm_with_tools это мозг операции, а ToolNode - исполнительные миньоны.

from langgraph.prebuilt import ToolNode

def compile_graph(self):
    tool_node = ToolNode(tools_list)

Создаём граф и накидываем на/в него наши узлы:

def compile_graph(self):
    tool_node = ToolNode(tools_list)
    workflow = StateGraph(StoryState)
    workflow.add_node('start', self.start_node)
    workflow.add_node('architect', self.architect_node)
    workflow.add_node('tools', tool_node)
    workflow.add_node('writer', self.writer_node)
    workflow.add_node('critic', self.critic_node)
    workflow.add_node('corrector', self.corrector_node)

Теперь к маршрутизации. У нас есть несколько прямых связей и два условных ребра (вызов инструментов и возврат черновика критиком). Для определения ребра с инструментами воспользуемся встроенной Langgraph функцией langgraph.prebuilt.tools_condition. Фактически она делает что-то похожее на то, что и мы, в узле Архитектора - смотрит, есть ли в последнем сообщении атрибут tool_calls и возвращает tools или __end__ Опишем грани:

def compile_graph(self):
    # код из предыдущего листинга
    workflow.set_entry_point('start')
    workflow.add_edge('start', 'architect')
    workflow.add_conditional_edges(
        'architect', tools_condition,
        {'tools': 'tools', '__end__': 'writer'})
    workflow.add_edge('tools', 'architect')
    workflow.add_edge('writer', 'critic')
    workflow.add_conditional_edges(
        'critic', self.route_after_critic,
        {'writer': 'writer', 'corrector': 'corrector'})
    workflow.set_finish_point('corrector')

На всякий случай - метод StateGraph.add_edge(a, b) создаст безусловный переход от узла с именем a к узлу с именем b, метод StateGraph.add_conditional_edges(entry, condition, {returned_val1: exit_1, returned_val2: exit_2, ...}) создаст условный переход с узла, с именем entry к одному из узлов (exit_1, exit_2) в зависимости от того, какое значение вернёт функция condition. StateGraph.set_entry_point и StateGraph.set_finish_point задают начало и конец графа соответственно.

Граф готов. Скомпилируем и вернём его для дальнейшего использования

def compile_graph(self):
  # код из предыдущих листингов
  return workflow.compile()

Запускаем. Ждём. Ждём. Ждём...

Настало время Ctrl+C, т.к. явно что-то не так. Идём в Langfuse и видим, что архитектор 26 раз вызывал инструмент. Либо ему не нравится ответ, либо где-то косяк.

Пересмотрев код узла Архитектора можно заметить, что в LLM не передаётся история сообщений. Таким образом у Архитектора каждый раз как первый раз. Исправляем. Для начала в web-интерфейсе Langfuse создадим новую версию промпта Архитектора, в которой просто удалим второе сообщение (которое пользователя). Теперь доработаем код:

    def architect_node(self, state: StoryState) -> dict[str, Any]:
        prompt = settings.langfuse.client.get_prompt(name='architect_prompt', label='latest').prompt[0]['content'])
        prompt_with_history = ChatPromptTemplate.from_messages([
            ('system', prompt),
            MessagesPlaceholder(variable_name='messages'),
        ])
        chain = prompt_with_history | self.llm_with_tools
        response = chain.invoke({'messages': state['messages']})
        if hasattr(response, 'tool_calls') and response.tool_calls:
            return {'messages': [response]}
        return {'messages': [response], 'task_for_writer': response.content}

Что поменялось: для начала мы получаем промпт не с помощью метода compile, а с помощью prompt. Данное свойство возвращает промпты в "сыром" виде (список словарей). Берём content первого элемента списка (наш System prompt) и получаем чистый текст промпта. Далее подмешиваем немного Langchain в наш Langgraph агент - создаём шаблон ChatPromptTemplate (фактически тот-же Langfuse чат с плейсхолдером) и собираем цепочку из шаблонизированного промпта и нашей LLM (с инструментами). Конечно, можно было заменить Langfuse промпте User секцию на Placeholder. Но Placeholder требует на вход список словарей вида {'role': , 'content': '} и нам пришлось бы мапить классы сообщений Langgraph в понятные Langfuse словарь, а это немного посложнее.

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

По трейсу видно, что Архитектор один раз вызвал инструменты, далее написал задание Писателю. Критику всё время что-то не нравилось - цикл отработал максимально возможное число раз. Можно сделать вывод, что Langfuse справился с трассировкой нашего относительно сложного графа.

Итоги

Связка LangGraph + LangFuse определённо стоит времени на настройку. Без трассировки тяжело было бы сразу понять, почему зависает Сказочный агент да и для отладки пришлось бы в каждый узел логирование вставлять. Ну и версионирование промптов, как по мне, очень удобная штука для настройки агента.

Код проекта доступен тут.

Буду рад обратной связи в комментариях.

Наблюдаемость LLM-агентов: Часть 2. Разработка и отладка графа