Function _find_tool_pair_indices

web/pgadmin/llm/compaction.py:185–219

Find indices of tool_call/tool_result groups that must stay together. An assistant message may contain multiple tool_calls, each with a corresponding tool result message. All messages in such a group must be dropped or kept together. Returns a mapping where every index in a group m

(
    messages: list[Message]
)

def _find_tool_pair_indices(
    messages: list[Message]
) -> dict[int, frozenset[int]]:
    """Find indices of tool_call/tool_result groups that must stay together.

    An assistant message may contain multiple tool_calls, each with a
    corresponding tool result message. All messages in such a group
    must be dropped or kept together.

    Returns a mapping where every index in a group maps to the full
    set of indices in that group.

    Args:
        messages: The message list.

    Returns:
        Dict mapping index -> frozenset of all indices in the group.
    """
    groups: dict[int, frozenset[int]] = {}

    for i, msg in enumerate(messages):
        if msg.role == Role.ASSISTANT and msg.tool_calls:
            tool_call_ids = {tc.id for tc in msg.tool_calls}
            group_indices = {i}
            for j in range(i + 1, len(messages)):
                if messages[j].role == Role.TOOL:
                    for tr in messages[j].tool_results:
                        if tr.tool_call_id in tool_call_ids:
                            group_indices.add(j)
                            break
            group = frozenset(group_indices)
            for idx in group:
                groups[idx] = group

    return groups


def compact_history(