Token Budget

Token Budget

TokenBudgetHook 는 Agent 실행의 토큰 사용량을 제한할 수 있게 해 주는 훅이에요. 설정한 임계값에 도달하면 Agent가 다음 LLM 호출 전에 멈춰요. 그때까지 수집된 메시지는 계속 사용할 수 있어요.

출처: 문서

본문

경고: TokenBudgetHook 은 실험적이에요. API는 일반적인 폐기(deprecation) 정책 없이 어떤 릴리스에서든 바뀔 수 있어요.

Agent 컴포넌트에 before_llm 훅 포인트로 등록되는 TokenBudgetHook 으로 구성돼요. 핵심 클래스는 TokenBudgetHook, 임포트 경로는 haystack.hooks.budget 이에요.

개요 (Overview)

토큰 예산은 Agent의 일반적인 훅 메커니즘의 한 용도예요. before_llm 에 등록된 TokenBudgetHook 은 매 LLM 호출 전에 Agent의 State 에 있는 누적 token_usage 를 설정된 임계값과 비교해요.

사용량이 임계값에 도달하거나 넘으면 훅이 stop_run 을 설정해요. Agent는 LLM 호출을 더 하지 않고 실행을 끝내고, exit_reason 을 "token_budget_exceeded" 로 설정해요.

예산이 적용되는 범위 (What the budget covers)

예산은 Agent의 채팅 생성기 응답에서 누적된 token_usage 에 적용돼요. 도구나 다른 훅이 한 호출은 집계되지 않아요. 예를 들어 도구가 자체 LLM 호출을 하면 그 호출이 쓴 토큰은 Agent의 예산에 계산되지 않아요.

훅은 매 LLM 호출 전에 사용량을 확인하므로, 총량이 임계값을 넘게 만든 호출은 이미 완료된 상태예요. 그래서 최종 사용량은 한 번의 LLM 호출 비용만큼 max_total_tokens 를 초과할 수 있어요.

사용법 (Usage)

기본 설정 (Basic setup)

before_llm 에 훅을 등록하고 max_total_tokens 로 토큰 임계값을 설정하세요. 아래의 임계값은 Agent가 보고서를 완성하기 전에 리서치 작업을 멈추기에 충분히 낮아요:

import random
from typing import Annotated
from haystack.components.agents import Agent
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage
from haystack.hooks.budget import TokenBudgetHook
from haystack.tools import tool

FACTS = [
    "Capybaras are the largest living rodents, weighing up to 65 kg. ",
    "Capybaras are highly social and live in groups of ten to twenty. ",
    "Capybaras are excellent swimmers and can stay underwater for five minutes. ",
    "Capybaras are famously relaxed and often share space with birds and monkeys. ",
]

@tool
def search(query: Annotated[str, "The search query"]) -> str:
    """Search the web."""
    # Placeholder: would call a real search API
    # Repeat the result to simulate a longer search response
    return random.choice(FACTS) * 20

agent = Agent(
    chat_generator=OpenAIChatGenerator(model="gpt-5-mini"),
    tools=[search],
    system_prompt="You are a research assistant. Search one aspect at a time before answering.",
    hooks={"before_llm": [TokenBudgetHook(max_total_tokens=3_000)]},
)
agent.warm_up()
result = agent.run(
    messages=[
        ChatMessage.from_user(
            "Research capybaras: size, social life, swimming and temperament."
        )
    ]
)
print(result["exit_reason"])
# >> token_budget_exceeded

Agent는 리서치 중간에 멈추고 지금까지 수집된 메시지를 보존해요. 더 높은 임계값으로 설정하면 보고서를 완성하고 exit_reason 으로 "text" 를 반환할 수 있어요.

마지막 메시지 추가 (Adding a final message)

예산이 실행을 멈추면 마지막 메시지가 최종 답변이 아니라 도구 결과일 수 있어요. add_final_message=True 로 설정하면 실행이 왜 끝났는지 설명하는 어시스턴트 메시지를 추가해요. 이 메시지가 last_message 가 돼요:

TokenBudgetHook(max_total_tokens=3_000, add_final_message=True)

메시지를 커스터마이즈하거나 max_agent_steps 같은 다른 종료 이유를 처리하려면 after_run 훅을 사용하세요. after_run 은 Agent가 종료 이유와 무관하게 끝난 뒤 실행돼요:

from haystack.components.agents.state import State
from haystack.dataclasses import ChatMessage
from haystack.hooks import hook

@hook
def explain_stop(state: State) -> None:
    if state.get("exit_reason") == "token_budget_exceeded":
        state.set(
            "messages",
            [ChatMessage.from_assistant("I ran out of budget before finishing.")],
        )

더 알아보기 (Learn more)