Как конечные автоматы помогают сделать агента надежнее и при чем тут pydantic-graph?

Когда пишешь библиотеку, рано или поздно упираешься в движок. Не в красивый внешний интерфейс и не в обёртки, а в ту часть внутри, которая гоняет процесс по состояниям: что-то сгенерировал, проверил, решил, что делать дальше, повторил. Пара флагов, цикл while, большой if посередине, и через месяц вы уже сами не помните, какие переходы там вообще возможны и почему одна из веток недостижима.

Недавно я собирал ровно такой движок и наткнулся на библиотеку, которая делает эту работу заметно аккуратнее. Называется pydantic-graph. Про неё почти не пишут, хотя на ней стоит весь pydantic-ai, агентский фреймворк от авторов Pydantic. Дальше я расскажу про неё на конкретном примере, харнессе надёжности для слабых языковых моделей.

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

Проблема: модель умеет, но ошибается по мелочи

Пример, на котором я всё обкатывал, такой:

Берём небольшую модель (у меня это Llama 3.1 8B Instruct) и просим решать задачи из HumanEval+ (кодовый бенчмарк для llm). Малоресурсные модели на таких задачах не то чтобы не справляются, они справляются через раз. С наскоку, без всякой обвязки, эта модель набирает pass@1 около 0.61. Грубо говоря, примерно две попытки из пяти уходят в мусор, хотя решение обычно почти правильное и спотыкается на ерунде.

Я прогнал по 20 сэмплов на каждую задачу из 164 и заглянул в лог ошибок. Вот характерные примеры:

• NameError: name 'encode_cyclic' is not defined. Модель вызвала вспомогательную функцию, которую сама же забыла дописать;

• max() arg is an empty sequence. Не обработан пустой вход, и на нём решение падает;

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

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

Полная таблица метрик первого прогона

Те же 164 задачи, 20 сэмплов на задачу, temp=0.2. Цифры по скрытым тестам (plus), в скобках по публичным (base). Кривая pass@k (число попыток):

k	plus	base
pass@1	0.607	0.665
pass@2	0.654	0.712
pass@5	0.698	0.758
pass@10	0.722	0.782

Идея: движок ведёт агента

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

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

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

pydantic-graph: граф состояний на типах

Ставится отдельно, тащить за собой весь pydantic-ai не нужно:

uv add pydantic-graph==1.105.0  # версия на момент написания статьи

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

Сама библиотека про языковые модели знать не обязана. Да, на ней собран pydantic-ai, но в основе это автомат общего назначения. Он подойдёт любому процессу с состояниями: пайплайну обработки заказа, машине состояний в игре, ретраям при походе во внешний сервис. Меня она подкупила тем, что читается как обычный Python без скрытой магии, а граф при этом остаётся проверяемым.

Контракт агента: промпт на входе, строка на выходе

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

@runtime_checkableclass AgentFn(Protocol):    """Один ход агента: промпт внутрь, сырой текст наружу.    Под этот контракт подгоняется и одиночная модель, и целая мультиагентная система."""    def __call__(self, prompt: str) -> str: ...

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

А можно на всю команды целиком. Планировщик, исполнители, вся внутренняя оркестрация прячутся внутрь, и снаружи харнесс по-прежнему видит ровно один агент: на вход промпт, на выходе строка. Движку всё равно, что там за рамкой.

Устройство графа

Рабочие узлы возвращаются в центральный узел решения (Decide):

Вот как один узел выглядит в коде. Обратите внимание на сигнатуру метода run:

@dataclassclass Generate(BaseNode[State, Deps, Result]):    """Узел с моделью. Качество ответа зависит от того, что repair-нода положила в промпт."""    async def run(self, ctx: GraphRunContext[State, Deps]) -> Validate:        state = ctx.state        prompt = f"{state.task}\n\n{state.hint}" if state.hint else state.task        out = ctx.deps.agent(prompt)  # тот самый агент из контракта        state.history.append(Attempt(prompt=prompt, output=out))        return Validate()  # переход выражен типом возврата

Вот это -> Validate и есть та деталь, ради которой я остался на библиотеке. Это не комментарий и не строчка в конфиге, а настоящая аннотация типа. Тайпчекер знает, что из узла генерации можно попасть только в узел проверки. Если я по ошибке верну узел, которого в этой точке быть не должно, мне подсветят это в IDE ещё до запуска.

Собирается всё билдером :

gb = GraphBuilder(state_type=State, deps_type=Deps, output_type=Result)gb.add(    g.node(Generate),    g.node(Validate),    g.node(Decide),    g.node(Repair),    g.node(Fallback))graph = gb.build()

Коротко про роли узлов. Узел генерации это единственное место с моделью. Узел проверки детерминированный, он сводит грязную семантику ответа к простому вердикту: прошло или нет, какого рода ошибка, короткая причина, доля пройденных проверок. Решающий узел (Decide) это чистая логика над этим вердиктом, историей попыток и остатком бюджета, без модели внутри. Repair-нода собирает следующий промпт из последней ошибки и снова зовёт генерацию. Узел отката срабатывает, когда попытки кончились: возвращает лучший по оценке вариант и помечает результат как деградировавший.

Валидаторы: сменные проверки в одном слоте

Узел проверки (Validate) не знает заранее, что считать правильным ответом. Он принимает любой коллбэк вида «строка на вход, вердикт на выход», и это превращается в слот, куда можно вставлять что угодно.

Здесь есть где проявить фантазию. Самые дешёвые проверки это механические. Например, разбор парсером, линтер, проверка типов, валидация JSON по схеме, банальное «а оно вообще запускается?». Они универсальные, работают на любой строке, но и сигнал дают слабый, ловят только формат и падения. Думаю, что сильный сигнал всегда доменный (для кода это прогон тестов). Сверху можно повесить модель в роли судьи, по одному судье или сразу ансамблем.

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

Repair-узел и куда его можно развивать

Сейчас Repair-узел у меня довольно простой: вот ошибка, перегенерируй с её учётом фидбэка. Своего рода промпт-инжиниринг. Но способы починки можно заложить разные. Например, делать роутинг по типу ошибки — упало в рантайме, пробуем модель побольше. Или прогонять несколько моделей, переданных пользователем, и брать лучшего кандидата.

Что получилось в итоге?

Харнесс поднимает pass@1 с 0.607 до 0.623. Абсолютная погрешность каждой клетки около ±0.034, и на первый взгляд прирост растворяется. Но оба режима гонялись на одних и тех же 164 задачах, поэтому смотреть надо на парную разницу по задаче: она составляет +0.016 ± 0.007, то есть устойчиво положительная, хоть и небольшая.

Что мне понравилось по ходу дела, помимо самих чисел. Всё это оказалось расширяемым и строго типизированным. Граф читается прямо в IDE, переходы нельзя перепутать незаметно.

Вывод

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

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

Но если процесс ветвится, имеет состояния и вам важно, чтобы IDE держала его топологию за вас, библиотека ложится хорошо. И, повторю, она вовсе не обязана быть про языковые модели.

Весь код оформлен в библиотеку и доступен на GitHub.