Coverage for haystack/hooks/budget/hooks.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7from haystack import logging 

8from haystack.components.agents.state.state import State 

9from haystack.components.agents.utils import _INPUT_TOKEN_KEYS, _OUTPUT_TOKEN_KEYS, _first_numeric 

10from haystack.core.serialization import default_from_dict, default_to_dict 

11from haystack.dataclasses import ChatMessage 

12from haystack.utils.experimental import _experimental 

13 

14logger = logging.getLogger(__name__) 

15 

16 

17_FINAL_MESSAGE_TEXT = "The Agent stopped because the token budget was exceeded." 

18 

19 

20@_experimental 

21class TokenBudgetHook: 

22 """ 

23 Stop an Agent run when its token usage reaches a configured budget. 

24 

25 The hook runs at the `before_llm` hook point and checks the cumulative token usage recorded in the Agent state. 

26 When the budget is reached, the run ends before the next LLM call with the exit reason `"token_budget_exceeded"`. 

27 

28 Only calls made by the Agent's chat generator contribute to `token_usage`; calls made by tools or other hooks are 

29 not included. 

30 

31 <!-- test-ignore --> 

32 ```python 

33 from haystack.components.agents import Agent 

34 from haystack.components.generators.chat import OpenAIChatGenerator 

35 from haystack.hooks.budget import TokenBudgetHook 

36 

37 agent = Agent( 

38 chat_generator=OpenAIChatGenerator(), 

39 tools=[web_search], 

40 hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100_000)]}, 

41 ) 

42 

43 result = agent.run(messages=[...]) 

44 ``` 

45 """ 

46 

47 allowed_hook_points = ("before_llm",) 

48 

49 def __init__(self, *, max_total_tokens: int, add_final_message: bool = False) -> None: 

50 """ 

51 Create a token budget hook. 

52 

53 :param max_total_tokens: Maximum cumulative token usage before the Agent is stopped. 

54 :param add_final_message: Whether to append an assistant message explaining why the Agent stopped. 

55 :raises ValueError: If `max_total_tokens` is less than 1. 

56 """ 

57 if max_total_tokens < 1: 

58 raise ValueError(f"`max_total_tokens` must be a positive number of tokens, got {max_total_tokens}.") 

59 self.max_total_tokens = max_total_tokens 

60 self.add_final_message = add_final_message 

61 

62 def run(self, state: State) -> None: 

63 """ 

64 Stop the Agent if its cumulative token usage has reached the budget. 

65 

66 :param state: Agent state containing the cumulative token usage. 

67 """ 

68 usage = state.data.get("token_usage") or {} 

69 # Not every chat generator reports `total_tokens`, so fall back to summing the input and output keys across 

70 # the known naming conventions. 

71 total_tokens = _first_numeric(usage, ("total_tokens",)) 

72 if not total_tokens: 

73 total_tokens = _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS) 

74 if total_tokens >= self.max_total_tokens: 

75 logger.warning( 

76 "Agent reached its token budget of {max_total_tokens} ({total_tokens} used); requesting a stop.", 

77 max_total_tokens=self.max_total_tokens, 

78 total_tokens=total_tokens, 

79 ) 

80 state.set("stop_run", "token_budget_exceeded") 

81 if self.add_final_message: 

82 state.set("messages", [ChatMessage.from_assistant(_FINAL_MESSAGE_TEXT)]) 

83 

84 def to_dict(self) -> dict[str, Any]: 

85 """ 

86 Serialize this hook to a dictionary. 

87 

88 :returns: Serialized representation of the hook. 

89 """ 

90 return default_to_dict(self, max_total_tokens=self.max_total_tokens, add_final_message=self.add_final_message) 

91 

92 @classmethod 

93 def from_dict(cls, data: dict[str, Any]) -> "TokenBudgetHook": 

94 """ 

95 Create a hook from its serialized representation. 

96 

97 :param data: Serialized hook data. 

98 :returns: The deserialized hook. 

99 """ 

100 return default_from_dict(cls, data=data)