Coverage for haystack/hooks/compaction/utils.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from haystack.dataclasses import ChatMessage, ChatRole 

6from haystack.token_counters import TokenCounter 

7from haystack.tools import ToolsType 

8 

9# Meta key marking a message that a compactor produced. Its value records which strategy ran. 

10_COMPACTION_META_KEY = "context_compaction" 

11 

12 

13def _leading_system_end(messages: list[ChatMessage]) -> int: 

14 """Return the end of the leading system-message block, excluding system messages created by compaction.""" 

15 for index, message in enumerate(messages): 

16 if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: 

17 return index 

18 return len(messages) 

19 

20 

21def _latest_user_index(messages: list[ChatMessage]) -> int | None: 

22 """Return the latest user message not produced by compaction.""" 

23 for index in reversed(range(len(messages))): 

24 message = messages[index] 

25 if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: 

26 return index 

27 return None 

28 

29 

30def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: 

31 """Return the messages at the given indices, in the order the indices are given.""" 

32 return [messages[index] for index in indices] 

33 

34 

35def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: 

36 """Return the messages the given indices leave out, in conversation order.""" 

37 left_out = set(indices) 

38 return [message for index, message in enumerate(messages) if index not in left_out] 

39 

40 

41def _is_compaction_message(message: ChatMessage, strategy: str, role: ChatRole | None = None) -> bool: 

42 """Return whether a message was produced by a compaction strategy and optionally has the requested role.""" 

43 marker = message.meta.get(_COMPACTION_META_KEY) 

44 has_role = role is None or message.is_from(role=role) 

45 return has_role and isinstance(marker, dict) and marker.get("strategy") == strategy 

46 

47 

48def _last_assistant_index(messages: list[ChatMessage]) -> int: 

49 """Return the index of the last assistant message, or -1 if none exists.""" 

50 for index in reversed(range(len(messages))): 

51 if messages[index].is_from(role=ChatRole.ASSISTANT): 

52 return index 

53 return -1 

54 

55 

56def _agent_step_spans(messages: list[ChatMessage], start: int) -> list[tuple[int, int]]: 

57 """ 

58 Return spans containing an assistant message and all immediately following tool results. 

59 

60 :param messages: The conversation to analyze, oldest to newest. 

61 :param start: The index to start searching for steps from. We recommend starting after the latest user message, or 

62 after the leading system messages when there is no user message. Otherwise, the returned spans may include 

63 steps that are not part of the current task. 

64 :returns: A list of spans, where each span is a tuple of (start_index, end_index) and end_index is exclusive. 

65 """ 

66 # e.g. a span of (2, 5) means messages[2:5] are part of the same step 

67 spans: list[tuple[int, int]] = [] 

68 index = start 

69 while index < len(messages): 

70 # Only assistant messages can start a step; skip any other messages. 

71 if not messages[index].is_from(role=ChatRole.ASSISTANT): 

72 index += 1 

73 continue 

74 # Find the end of the step by looking for the first message that is not a tool result. 

75 end = index + 1 

76 while end < len(messages) and messages[end].tool_call_results: 

77 end += 1 

78 # Record the span and continue searching for the next step. 

79 spans.append((index, end)) 

80 # Reset the index to the end of the current step to avoid overlapping spans. 

81 index = end 

82 return spans 

83 

84 

85def _current_agent_step_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: 

86 """ 

87 Return message-index groups for complete Agent steps belonging to the current task. 

88 

89 :param messages: The conversation to analyze, ordered oldest to newest. 

90 :param system_end: The end of the leading system-message block. Steps are looked for from here when the 

91 conversation has no user message to anchor on. 

92 :param task_index: The index of the user message anchoring the current task, or None when there is none. Steps are 

93 looked for from the message after it. 

94 :returns: One group of message indices per step, ordered oldest step first. A step is an assistant message and all 

95 immediately following tool results, so a group holds a tool call together with its results. 

96 """ 

97 step_start = task_index + 1 if task_index is not None else system_end 

98 return [list(range(start, end)) for start, end in _agent_step_spans(messages=messages, start=step_start)] 

99 

100 

101def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: 

102 """ 

103 Return spans for complete real-user turns in a bounded section of conversation history. 

104 

105 Each turn begins with a user message not created by compaction and continues up to the next such message. 

106 """ 

107 user_indices = [ 

108 index 

109 for index in range(start, end) 

110 if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta 

111 ] 

112 return [ 

113 (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) 

114 for position, index in enumerate(user_indices) 

115 ] 

116 

117 

118def _historical_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: 

119 """ 

120 Return message-index groups for complete historical turns preceding the current task. 

121 

122 :param messages: The conversation to analyze, ordered oldest to newest. 

123 :param system_end: The end of the leading system-message block, where the history begins. 

124 :param task_index: The index of the user message anchoring the current task, where the history ends. None when 

125 there is no such message, in which case the history is empty and everything belongs to the current task. 

126 :returns: One group of message indices per turn, ordered oldest turn first. A turn is a user message that 

127 compaction did not produce, together with everything that follows it up to the next one. 

128 """ 

129 historical_end = task_index if task_index is not None else system_end 

130 return [ 

131 list(range(start, end)) 

132 for start, end in _historical_turn_spans(messages=messages, start=system_end, end=historical_end) 

133 ] 

134 

135 

136def _estimated_context_tokens( 

137 messages: list[ChatMessage], context_tokens: int, token_counter: TokenCounter, tools: ToolsType | None = None 

138) -> int: 

139 """ 

140 Estimate the size of the whole conversation. 

141 

142 :param messages: The conversation, oldest to newest. 

143 :param context_tokens: The `context_tokens` state key which is computed using the provider's own token counting. 

144 :param token_counter: The counter to measure the unaccounted messages with. 

145 :param tools: Tools whose schemas are sent alongside the messages. These are counted when provider usage is absent. 

146 :returns: The estimated total token count. 

147 """ 

148 # Nothing sent yet, or a generator that reports no usage, so count everything. 

149 if context_tokens == 0: 

150 return token_counter.count(messages=messages, tools=tools) 

151 # Only need to estimate the tool result messages after the last assistant message 

152 tool_result_messages = messages[_last_assistant_index(messages=messages) + 1 :] 

153 return context_tokens + token_counter.count(messages=tool_result_messages)