Coverage for haystack/token_counters/types/protocol.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, Protocol 

6 

7from haystack.core.serialization import default_from_dict 

8from haystack.dataclasses import ChatMessage 

9from haystack.tools import ToolsType 

10 

11 

12class TokenCounter(Protocol): 

13 """ 

14 Estimates the number tokens used by a list of messages. 

15 

16 Implement `to_dict` so the counter's settings survive serialization. The default `from_dict` passes them straight 

17 back to the constructor, which is enough for plain values; override it when `to_dict` emitted something that has to 

18 be rebuilt first, such as a `Secret` or a nested component. 

19 """ 

20 

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

22 """ 

23 Return the estimated number of tokens in the given messages. 

24 

25 :param messages: The messages to measure. 

26 :param tools: Tools whose schemas are sent alongside the messages, and so consume tokens too. Pass them to have 

27 them counted; leave as None to measure the messages alone. 

28 :returns: The estimated token count. 

29 """ 

30 ... 

31 

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

33 """Serialize the counter to a dictionary.""" 

34 ... 

35 

36 @classmethod 

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

38 """Deserialize the counter from a dictionary.""" 

39 return default_from_dict(cls, data)