ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить token budget для агента

1. Цель задачи

Научиться контролировать расход токенов в LLM-агенте, чтобы предотвратить неконтролируемые затраты (runaway costs). Реализовать механизм budget: установить лимит 5000 токенов на сессию агента и принудительно останавливать его работу при превышении. В процессе освоить методы мониторинга потребления токенов, интеграцию лимита в цикл агента]] и обработку граничных случаев.

Ключевой результат Агент гарантированно завершает сессию при достижении суммарного потребления токенов (вход + выход) 5000 токенов, а логи фиксируют точное потребление, момент остановки и причину. Стоимость сессии никогда не превысит заданный лимит.

2. Исходные данные

Если нет реального агента — симулируем:

1. Создаём простого ReAct-агента на langgraph или smolagents с 2–3 инструментами (калькулятор, поиск по Wikipedia, чтение файла).
2. Подключаем LLM через OpenAI API или локальную модель (через Ollama / vLLM) — для отслеживания токенов важен любой подсчёт.
3. Добавляем глобальный счётчик токенов, который суммирует input_tokens + output_tokens каждого вызова LLM.

Если API не возвращает число токенов (например, Ollama), используем tiktoken для подсчёта по строке запроса/ответа.

3. Технологический стек

4. Этапы выполнения

Этап 1: Подготовка окружения и базовый подсчёт токенов (30 минут)

Действия

1. Установить зависимости:

   pip install langchain langgraph openai tiktoken python-dotenv pytest
   

2. Создать файл token_counter.py с функцией подсчёта токенов для строки:

   import tiktoken
   
   def count_tokens(text: str, model: str = "gpt-4") -> int:
       encoder = tiktoken.encoding_for_model(model)
       return len(encoder.encode(text))
   

3. Написать тест проверки подсчёта:

   def test_count_tokens():
       assert count_tokens("Hello, world") == 4  # зависит от модели
   

4. Подготовить агента: если его нет, создать минимального агента с одним инструментом (например, calculator). Убедиться, что каждый вызов LLM возвращает usage (OpenAI) или можно посчитать строки.

Ожидаемый результат этапа Работает функция count_tokens, агент способен выполнять простые запросы.

Этап 2: Внедрение глобального аккумулятора токенов (60 минут)

Действия

1. Добавить в агента класс TokenTracker:

   from dataclasses import dataclass, field
   
   @dataclass
   class TokenTracker:
       budget: int = 5000
       used: int = 0
       stopped: bool = False
   
       def consume(self, tokens: int) -> bool:
           """Возвращает True, если budget превышен"""
           self.used += tokens
           if self.used >= self.budget:
               self.stopped = True
           return self.stopped
   

2. Модифицировать функцию вызова LLM так, чтобы перед запросом и после ответа увеличивать счётчик:

   def llm_call(prompt: str, tracker: TokenTracker) -> str:
       if tracker.stopped:
           raise RuntimeError("Token budget exceeded")
       
       # Подсчёт токенов запроса
       input_tokens = count_tokens(prompt)
       tracker.consume(input_tokens)
       if tracker.stopped:
           raise RuntimeError("Token budget exceeded (input)")
   
       # Вызов LLM (пример OpenAI)
       response = client.chat.completions.create(...)
       output_tokens = response.usage.completion_tokens  # или count_tokens(response.choices[0].message.content)
       tracker.consume(output_tokens)
       
       if tracker.stopped:
           raise RuntimeError("Token budget exceeded (output)")
       
       return response.choices[0].message.content
   

3. Интегрировать TokenTracker в цикл агента (например, в AgentExecutor LangChain через callback или кастомный on_llm_start/on_llm_end).

Ожидаемый результат этапа Любой вызов LLM увеличивает счётчик; при превышении лимита выбрасывается исключение.

Этап 3: Остановка агента при превышении лимита (45 минут)

Действия

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

   def safe_run(agent, prompt: str, tracker: TokenTracker):
       try:
           result = agent.run(prompt, callbacks=[TrackerCallback(tracker)])
           return result
       except RuntimeError as e:
           log.warning(f"Session stopped: {e}")
           return {"error": str(e), "used_tokens": tracker.used}
   

2. В конце сессии (нормальной или прерванной) записать в лог:

   • session_id (например, UUID)
   • total_tokens_used
   • stopped (True/False)
   • budget

3. Тест: отправить агенту запрос, который гарантированно превысит лимит (например, генерация длинного текста или многошаговая цепочка). Убедиться, что агент не "зависает" и завершается с ошибкой.

Ожидаемый результат этапа Агент корректно останавливается при превышении budget, лог содержит ключевую информацию.

Этап 4: Тестирование сценариев и граничные случаи (30 минут)

Действия

1. Подготовить минимум 5 тестовых сценариев:

   • короткий запрос (должен пройти)
   • длинный запрос (превышение на входе)
   • короткий запрос + длинный ответ (превышение на выходе)
   • итеративный агент с несколькими вызовами (суммарное превышение)
   • запрос, требующий много шагов (проверка, что остановка происходит без потери данных)

2. Написать pytest-тесты:

   def test_short_query():
       tracker = TokenTracker(budget=5000)
       result = safe_run(agent, "Hello", tracker)
       assert result is not None
       assert result.get("error") is None
       assert tracker.used < 5000
   

3. Проверить, что после остановки агент не делает новых вызовов LLM (нет "траты токенов" после лимита).

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

Этап 5: Мониторинг и экспорт метрик (30 минут)

Действия

1. Создать скрипт monitor.py, который в реальном времени читает логи и выводит:

   Session 12345 | used: 1230 / 5000 | stopped: False
   Session 12346 | used: 5120 / 5000 | stopped: True  (OVERRUN!)
   

2. (Опционально) Экспорт в Prometheus с помощью prometheus_client:

   from prometheus_client import Gauge, start_http_server
   token_gauge = Gauge('agent_token_usage', 'Current token usage', ['session_id'])
   

3. Убедиться, что при превышении лимита в лог пишется сообщение с уровнем WARNING и указанием превышения.

Ожидаемый результат этапа Логи содержат метрики по каждой сессии, скрипт мониторинга показывает актуальное состояние.

5. Критерии приемки (Definition of Done)

• Агент не превышает 5000 токенов за сессию (сумма input + output).
• При достижении лимита сессия корректно останавливается (без lossy side-effects).
• Механизм не блокирует сессии, которые не достигают лимита.
• В логах каждой сессии присутствуют поля: session_id, total_tokens, budget, stopped, timestamp.
• Написан как минимум один тест, проверяющий превышение лимита.
• Написан тест, проверяющий, что нормальная сессия завершается успешно.
• Код покрыт комментариями и содержит README с инструкцией по запуску.
• Мониторинг (лог или Prometheus) показывает количество активных и остановленных сессий.
• Нет "утечки токенов" (token leak) — счётчик не сбрасывается раньше времени.

6. Ожидаемый результат

Основной файл agent_with_budget.py (или аналогичный), содержащий:

• класс TokenTracker
• модифицированный агент с поддержкой budget
• функцию safe_run
• примеры вызова

Дополнительно

• tests/test_token_budget.py — набор тестов (минимум 5 тестов)
• monitor.py — скрипт мониторинга (опционально)
• requirements.txt — список зависимостей
• logs/ — пример записей лога с остановленной сессией

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: При первом выполнении (знакомство с фреймворком агентов) рекомендуется добавить 1 час резерва.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я подключил tiktoken и проверил, что функция подсчёта даёт корректные значения для используемой модели.
• Я добавил TokenTracker в цикл агента так, что каждый вызов LLM увеличивает счётчик.
• Я реализовал исключение RuntimeError при превышении лимита и убедился, что агент не делает новых вызовов.
• Я написал минимум 5 тестов, включая тест на превышение и на нормальное выполнение.
• Я проверил, что логи содержат все обязательные поля (session_id, used, stopped).
• Я протестировал асинхронный сценарий (если агент async) и убедился, что блокировки работают.
• Я описал в README как запустить агента и как интерпретировать логи.