Coverage for haystack/token_counters/tiktoken_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.lazy_imports import LazyImport 

10from haystack.token_counters.types import TokenCounter 

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

12from haystack.tools import ToolsType 

13 

14with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports: 

15 import tiktoken 

16 

17 

18class TiktokenCounter(TokenCounter): 

19 """ 

20 Counts tokens locally with `tiktoken`, OpenAI's byte-pair encoder. 

21 

22 Counting is an estimate, and two limits are worth knowing before relying on it: 

23 - **It is text-only**, so images and files get the flat `tokens_per_image` / `tokens_per_file` estimate rather 

24 than a real count. 

25 - **It is OpenAI's encoder.** Other providers tokenize differently, so expect the count to drift on them. 

26 

27 ## Usage Example: 

28 ```python 

29 from haystack.dataclasses import ChatMessage 

30 from haystack.token_counters import TiktokenCounter 

31 

32 counter = TiktokenCounter(encoding="o200k_base") 

33 messages = [ 

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

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

36 ] 

37 token_count = counter.count(messages) 

38 print(f"Token count: {token_count}") 

39 ``` 

40 """ 

41 

42 def __init__(self, encoding: str = "o200k_base", tokens_per_image: int = 85, tokens_per_file: int = 1000) -> None: 

43 """ 

44 Initialize the counter. 

45 

46 :param encoding: The `tiktoken` encoding to count with. The default, `o200k_base`, is what current OpenAI 

47 models use. 

48 :param tokens_per_image: Tokens to charge per image, which the tokenizer cannot measure. The default is what 

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

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

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

52 :raises ImportError: If `tiktoken` is not installed. 

53 """ 

54 tiktoken_imports.check() 

55 self.encoding = encoding 

56 self.tokens_per_image = tokens_per_image 

57 self.tokens_per_file = tokens_per_file 

58 self._encoder: Any = None 

59 

60 def warm_up(self) -> None: 

61 """Load the encoder, downloading its vocabulary if it is not already cached.""" 

62 if self._encoder is not None: 

63 return 

64 self._encoder = tiktoken.get_encoding(self.encoding) 

65 

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

67 """ 

68 Return the estimated number of tokens used by the given messages. 

69 

70 :param messages: The messages to measure. 

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

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

73 """ 

74 if not messages and not tools: 

75 return 0 

76 self.warm_up() 

77 text_tokens = len(self._encoder.encode(_rendered_conversation(messages) + _rendered_tools(tools))) 

78 return text_tokens + _non_text_tokens( 

79 messages, tokens_per_image=self.tokens_per_image, tokens_per_file=self.tokens_per_file 

80 ) 

81 

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

83 """ 

84 Serialize the counter. 

85 

86 :returns: A dictionary representation of the counter. 

87 """ 

88 return default_to_dict( 

89 self, encoding=self.encoding, tokens_per_image=self.tokens_per_image, tokens_per_file=self.tokens_per_file 

90 )