Coverage for haystack/token_counters/utils.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import json 

6from collections.abc import Callable 

7 

8from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent 

9from haystack.dataclasses.chat_message import ChatMessageContentT, ToolCallResultContentT 

10from haystack.tools import ToolsType, flatten_tools_or_toolsets 

11 

12# Builds the stand-in for message content that has no text form, such as an image. 

13_PlaceholderFn = Callable[[ChatMessageContentT], str] 

14 

15 

16def _non_text_placeholder(content: ChatMessageContentT) -> str: 

17 """A short stand-in, such as `<image>`, for message content that has no text form.""" 

18 if isinstance(content, ImageContent): 

19 return "<image>" 

20 if isinstance(content, FileContent): 

21 return f"<file: {content.filename or 'unnamed'}>" 

22 return f"<{type(content).__name__}>" 

23 

24 

25def _tool_result_text(result: ToolCallResultContentT, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: 

26 """A tool result as a single string, with placeholders standing in for any non-text parts.""" 

27 if isinstance(result, str): 

28 return result 

29 return "".join(block.text if isinstance(block, TextContent) else placeholder(block) for block in result) 

30 

31 

32def _render_message( 

33 message: ChatMessage, placeholder: _PlaceholderFn = _non_text_placeholder, include_tool_call_ids: bool = False 

34) -> str: 

35 """ 

36 One message as one or more lines of plain text. 

37 

38 Reasoning content is deliberately left out: providers discard it between turns, so it is not part of the context 

39 being measured. 

40 

41 :param message: The message to render. 

42 :param placeholder: Builds the stand-in for content that has no text form. The default is short and stable, which 

43 is what a counter needs because the stand-in's own length is what gets measured. A caller rendering for a model 

44 to read can pass one that describes the content instead. 

45 :param include_tool_call_ids: Whether to include tool call IDs in the rendered message. 

46 :returns: The rendered message. 

47 """ 

48 role = message.role.value 

49 # A tool-result msg only carries tool_call_results, so it is rendered on its own and labelled with the tool that 

50 # produced it. 

51 if results := message.tool_call_results: 

52 return "\n".join( 

53 f"[tool:{result.origin.tool_name}" 

54 f"{' id=' + result.origin.id if include_tool_call_ids and result.origin.id else ''}" 

55 f"{' (error)' if result.error else ''}] " 

56 f"{_tool_result_text(result.result, placeholder=placeholder)}" 

57 for result in results 

58 ) 

59 

60 lines: list[str] = [] 

61 if texts := message.texts: 

62 lines.append(f"[{role}] " + "\n".join(texts)) 

63 for call in message.tool_calls: 

64 arguments = json.dumps(call.arguments, default=str, sort_keys=True) 

65 call_id = f" id={call.id}" if include_tool_call_ids and call.id else "" 

66 lines.append(f"[{role} -> tool_call{call_id}] {call.tool_name}({arguments})") 

67 # Images and files cost tokens too, so they need a stand-in rather than being skipped. 

68 non_text: list[ChatMessageContentT] = [*message.images, *message.files] 

69 for content in non_text: 

70 lines.append(f"[{role}] {placeholder(content)}") 

71 return "\n".join(lines) if lines else f"[{role}] <no content>" 

72 

73 

74def _rendered_conversation( 

75 messages: list[ChatMessage], 

76 *, 

77 placeholder: _PlaceholderFn = _non_text_placeholder, 

78 include_tool_call_ids: bool = False, 

79) -> str: 

80 """The whole conversation as one plain-text block, which is what a counter measures.""" 

81 return "\n".join( 

82 _render_message(message, placeholder=placeholder, include_tool_call_ids=include_tool_call_ids) 

83 for message in messages 

84 ) 

85 

86 

87def _rendered_tools(tools: ToolsType | None) -> str: 

88 """ 

89 Tool schemas as one JSON block, standing in for what a provider sends alongside the messages. 

90 

91 :param tools: The tools to render, or None. 

92 :returns: The rendered schemas, or an empty string when there are none. 

93 """ 

94 if not tools: 

95 return "" 

96 return json.dumps([tool.tool_spec for tool in flatten_tools_or_toolsets(tools)], default=str, sort_keys=True) 

97 

98 

99def _non_text_tokens(messages: list[ChatMessage], *, tokens_per_image: int, tokens_per_file: int) -> int: 

100 """ 

101 A flat estimate for the images and files in a conversation, which have no text to measure. 

102 

103 :param messages: The conversation to measure. 

104 :param tokens_per_image: Tokens to charge per image. 

105 :param tokens_per_file: Tokens to charge per file. 

106 :returns: The estimated token count for the non-text content. 

107 """ 

108 images = 0 

109 files = 0 

110 for message in messages: 

111 images += len(message.images) 

112 files += len(message.files) 

113 # Images and files returned by a tool are nested inside the tool result 

114 for result in message.tool_call_results: 

115 if isinstance(result.result, str): 

116 continue 

117 images += sum(isinstance(block, ImageContent) for block in result.result) 

118 files += sum(isinstance(block, FileContent) for block in result.result) 

119 return images * tokens_per_image + files * tokens_per_file