Приветствую читателей.
Мы пытались построить LLM-чат для продакшена.
Через месяц у нас был 20k-токенный prompt, 50 тулзов и ответы по 2 минуты.
В итоге пришлось отказаться от ReAct и перейти на LLMCompiler.

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

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

Простейший пример чата

history = []
system_prompt = """Keep the conversation going"""
history.append(system_prompt)

while True:
	user_message = input("Say something to llm")
	history.append(user_message)
	ai_message = llm.request(history)
	show_message_to_user(ai_message)
	history.append(ai_message)

Дальше появилась идея дать возможность ллм использовать тулзы, работает это так, что на вход вы добавляете строчку в духе: у тебя есть тулза get_weather, если пользователь спрашивает погоду, верни только название тулзы, таким образом ллм возвращает вам название тулзы, вы возвращаете в ллм ответ от тулзы и теперь у нее есть данные из внешнего апи, после этого ллм может. удовлетворить ответ пользователя.

Простейший прмер реализации агента с тулзами (ReACT agent)

history = []

def get_weather() -> str: ...

system_prompt = """
Keep the conversation going.
If user asks for a weather, answer with `get_weather` or read a response from this tool if available.
"""

history.append(system_prompt)

while True:
	user_message = input("Say something to llm")
	history.append(user_message)
	ai_message = llm.request(history)
	
	if ai_message == "get_weather":
	    tool_response = get_weather()
	    tool_response = f"get_weather response is: {tool_response}"
	    history.append(tool_response)
	    
	ai_message = llm.request(history) # llm видит в истории чата ответ от get_weather, вызываем еще раз ллм чтобы сформировать конечный ответ от ллм а не от тулзы
	
	show_message_to_user(ai_message)
	history.append(ai_message)

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

Шаблоны того как ллм должна форматировать свои ответы
Интерпретации намерений пользователей, например, "last week" это прошлая законченая неделя или 7 дней включая сегодня
Запрет на упоминание конкурентов
запрет на нецелевые запросы, например, «как создать ядерное оружие?»
как отвечать пользователю на те или иные ошибки
интеграция профиля пользователя, чат должен знать базовую информацию о пользователе
Ускорение ответов от чата
Запрет на фабрикацию данных от ллм

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

На самом деле это не главная беда, тут же начали вылезать основные минусы которые связанны с данной архитектурой

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

Одновременно с этим по из-за недостатка времени и опыта мы наделали следуйщих ошибок:

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

Времени было не много и нам очень хотелось чтобы чат перестал фабриковать данные и начал отвечать быстрее, переписывать 50 тулзов заняло бы слишком много времени

Выбор пал на архитектуру LLMCompiler, в своей сути она состоит из 2 агентов, я добавил третьего
Первый агент планирует вызов тулзов, строит цепочки и резолвит куда какие аргументы передать
Второй агент получает на вход то что сделал планировщик, историю чата и тулзы и решает достаточно ли этих данных чтобы удовлетворить запрос пользователя, если нет, дает инструкцию почему и отправляет ее об��атно в планировщик, если да, отправляет третьему агенту
Третий агент - composer, он получает историю чата, результаты вызовов тулзов и содержит в своем промпте информацию о шаблонах, какие данные как рендерить, это могут быть таблицы, параграфы, различные значки и так далее

Простейшая реализация

planner_prompt = """
Generate DAG of `tools` instances to satisfy latest user's request from `history`
<history>{history}</history> 
<tools>{tools}</tools>

example:
[{
	"id": "weat1",
	"name": "get_weather",
	"args": {"location": "London"},
	"dependencies": [],
	"thought": "I need to get weather"
},
{
	"id": "solv",
	"name": "solver",
	"args": {},
	"dependencies": ["weat1"],
	"thought": "All data is collected"
}
]
"""

solver_prompt = """
Check if the `tools` outputs has enough data to answer latest question from `history`, if it's enough, answer with `--answer--`

<history>{history}</history>
<tools>{tools}</tools>
"""

composer_prompt = """
	Using data from `tool_calls`, answer the latest question from `history`
	<history>{history}</history>
	<tool_calls>{plan}</tool_calls>
"""

plan_stream = await llm.request(planner_prompt)
messages = []
tools = [get_weather, get_location, solver]

plans = []

while True:
	tasks_pool = []
    user_input = input("Ask something")
    messages.append(user_input)
    plan_stream = llm.request(planner_prompt, history=messages, tools=tools, stream=True)
    async for task in plan_stream:
        schedule_task(task, tasks_pool)
        
    plan = await asyncio.gather(tasks_pool)
    plans.append(plan)
    
    solver_decision = await llm.request(solver_prompt, tools=tools, history=messages, tool_calls=plan)
    
    if "--answer--" != solver_decision:
	    plan.need_replan_reason = solver_decision
        continue
        
	ai_message = await llm.request(composer_prompt, tool_calls=plan, history=messages)
	show_message_to_user(ai_message)
	messages.append(ai_message)

Что стоит отметить:

Планировщик строит граф вызов тулзов, можно выбрать любой формат, мы выбрали json.
Ответ от планировщика стоит читать как стрим а не ждать полного ответа, мы сделали ответ в json формате поэтому мы могли парсить и запускать задачи на лету, в основном как только стрим заканчиваеться, все таски уже выполнены
Я добавил последнюю задачу всегда solver, в сути это способ планера общаться с солвером, он может поянить почему были сделаны те или иные решения поскольку каждый шаг имеет поле "thought" которое реализует паттер Chain of Thought, например, планировщик видит что пользователь нарушает какое-то правило поэтому в solver ноде он может сослаться на это, таким образом solver'у ноде будет проще сделать решение

Столкновение с реальностью и последствия кривого дизайна тулзов

Окей, мы реализовали этот паттерн, но как это насадить на то, что у нас уже есть? А именно, как построить цепочку вызовов если инпут и аутпут тулзов - json формат? мы не можем просто взять и передать аргументы одной функции в другую, например, get_weather(get_location()), эти тулзы могут иметь сложную структуру на входе и на выходе.

Первое что нужно было реализовать - язык разрешения аргументов, я использовал json pointer, это просто способ указать конкретный путь в json объекте, например, data/items/0/name
Таким образом, ллм может связывать вызовы между тулзами в своем графе, например,

[{
	"id": "loc1",
	"name": "get_location",
	"args": {},
	"thought": "...",
	"dependencies": []
},
{
	"id": "weather",
	"name": "get_weather",
	"args": {"location": "${loc1/data/value}}",
	"dependencies": ["loc1"],
	"thought": "..."
}
]

Как вы могли заметить, args теперь ссылается на прошлый вызов {"location": "${loc1/data/value}}, распарсить это в в коде не являеться проблемой.

Вторая проблема - несостыковка типов и форматов одной тулзы и другой, например, одна тулза может вернуть int а зависимая нода принимает то же значение но в типе str, также неясно что делать, например, с различными форматами времени

На ум пришла идея создать костыль, назвал я эту функцию transform_data(value: str, format_desc: str, output_type: type), внутри это был простой вызов в ллм, format_desc описывает как именно преобразовать value а затем результат приводился в питоновский тип

Позже мы все таки решили поправить наши тулзы и сделали одинаковые типы и добавили описание респонсов каждой тулзы чтобы у планировщика было достаточно информации чтобы строить граф без костылей, transform_data была удалена.

Проклятые списки

Казалось бы, удача, но, спустя несколько дней стало ясно что планировщик подозрительно часто делает реплан, стали разбираться и осознали что не учли важный момент, планировщик не может надежно построить план когда ему нужно делать ссылки на списки поскольку у него нет информации о том, сколько объектов в каждом списке на который он ссылаеться, например, пользователь может попросить покажи айди всех проектов для организации ABC, первый вызов будет get_projects_by_org_id() -> list[Project], следуйщий вызов должен быть по идее get_project_insights(project_id: str) -> Project но в рамках одного вызова llm планировщик не может знать сколько объектов вернет get_projects_by_org_id, решилась эта проблема прогревом данных, перед каждым запросом мы в контекст передаем иерархию всех объектов, в своей сути это json которые представляет собой дерево типа organization -> projects -> service -> instance, для маленьких организаций это работает отлично, для больших мы можем срезать service / instance чтобы уменьшить потребление контекста.

А что собственно изменилось

После внедрения llmcompiler'a изменилось следуйщее:

Даже имеея большой дам сырых данных ча почти перестал выдумывать метрики, если он чего-то не знает, говорит правду
Появился контроль над форматом респонса, компоузер агент содержит достаточно мало контекста, это позволяет детально описывать как должен выглядеть финальный ответ пользователю
Чат вызывает весь необходимый набор тулзов, перестал лениться
Достаточно просто находить упавшие запросы, у нас стоит рекунсия на 3 плана если за 3 плана солвер не согласился с набором тулзов планировщика, мы отдаем ответ пользователю в любом случае
Улучшеная валидация на уровне 3-х агентов, агенты имеют общие части промпта, в основном это касаеться всяких защитных промптов, если 1 агент проигнорировал инструкцию, другой агент поправит, если не поправил другой, поправит третий
Если планировщик справился с первого раза, скорость ответа увеличилась до 80% в зависимости он набора тулзов

Выводы и уроки

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

Если вам понравилась статья, подписывайтесь на телеграм

Почему ReAct-агенты ломаются в продакшене и чем их заменить