Сколько токенов съедает прогон агента

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

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

Задание

Реализуйте функцию total_prompt_tokens(trace). Аргумент — список пар (роль, размер) в том порядке, в каком сообщения появлялись в диалоге. Роль — одна из трёх строк: "user", "ai", "tool". Размер — длина сообщения в токенах.

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

Пример

trace = [("user", 10), ("ai", 5), ("tool", 3), ("ai", 7)]
total_prompt_tokens(trace)  # 28

Первый вызов модели видит только запрос пользователя — это 10 токенов. Второй видит запрос, ответ модели и результат инструмента: 10 + 5 + 3 = 18. Итого 28.

Примечание

Ответ модели и результаты инструментов в свой же вызов не попадают — они появляются в диалоге после него. Пустая трасса и трасса без единой роли "ai" дают 0.

Разрешена только стандартная библиотека Python 3.12. Решение должно обходить список один раз — не считайте сумму заново на каждом вызове модели.