Coverage for haystack/token_counters/approximate_counter.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.core.serialization import default_to_dict 

8from haystack.dataclasses import ChatMessage 

9from haystack.token_counters.types import TokenCounter 

10from haystack.token_counters.utils import _non_text_tokens, _rendered_conversation, _rendered_tools 

11from haystack.tools import ToolsType 

12 

13 

14class ApproximateTokenCounter(TokenCounter): 

15 """ 

16 Estimates tokens from text length using a flat ratio of characters to tokens. 

17 

18 ## Usage Example: 

19 ```python 

20 from haystack.dataclasses import ChatMessage 

21 from haystack.token_counters import ApproximateTokenCounter 

22 

23 counter = ApproximateTokenCounter(chars_per_token=4.0) 

24 messages = [ 

25 ChatMessage.from_user("Hello, how are you?"), 

26 ChatMessage.from_assistant("I'm good, thank you! How can I assist you today?") 

27 ] 

28 token_count = counter.count(messages) 

29 print(f"Estimated token count: {token_count}") 

30 ``` 

31 """ 

32 

33 def __init__(self, chars_per_token: float = 4.0, tokens_per_image: int = 85, tokens_per_file: int = 1000) -> None: 

34 """ 

35 Initialize the counter. 

36 

37 :param chars_per_token: How many characters to treat as one token. 

38 :param tokens_per_image: Tokens to charge per image, which has no text to measure. The default is what 

39 OpenAI charges for a small image; raise it if you send large ones. 

40 :param tokens_per_file: Tokens to charge per file. A rough stand-in for a short document, since the real 

41 cost depends on the page count; raise it if you send long ones. 

42 :raises ValueError: If `chars_per_token` is not positive. 

43 """ 

44 if chars_per_token <= 0: 

45 raise ValueError(f"`chars_per_token` must be greater than 0, got {chars_per_token}.") 

46 self.chars_per_token = chars_per_token 

47 self.tokens_per_image = tokens_per_image 

48 self.tokens_per_file = tokens_per_file 

49 

50 def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int: 

51 """ 

52 Return the estimated number of tokens the given messages occupy. 

53 

54 :param messages: The messages to measure. 

55 :param tools: Tools whose schemas are sent alongside the messages, and so consume tokens too. 

56 :returns: The estimated token count, or `0` when there is nothing to measure. 

57 """ 

58 if not messages and not tools: 

59 return 0 

60 text = _rendered_conversation(messages) + _rendered_tools(tools) 

61 text_tokens = int(len(text) / self.chars_per_token) 

62 return text_tokens + _non_text_tokens( 

63 messages=messages, tokens_per_image=self.tokens_per_image, tokens_per_file=self.tokens_per_file 

64 ) 

65 

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

67 """ 

68 Serialize the counter. 

69 

70 :returns: A dictionary representation of the counter. 

71 """ 

72 return default_to_dict( 

73 self, 

74 chars_per_token=self.chars_per_token, 

75 tokens_per_image=self.tokens_per_image, 

76 tokens_per_file=self.tokens_per_file, 

77 )