Class ContextPruner

java.lang.Object
uno.anahata.ai.context.pruning.ContextPruner
public class ContextPruner extends Object
Handles the pruning of messages and specific content parts from the conversation context.

This class implements complex logic to ensure that when a part is removed, all its dependent parts (e.g., the corresponding function call for a function response) are also identified and removed, maintaining the integrity of the conversation.

It supports manual pruning by ID, automatic pruning of old ephemeral tool calls, and stateful resource replacement.

  • Constructor Details

    • ContextPruner

      public ContextPruner(ContextManager contextManager)
      Constructs a new ContextPruner for the given ContextManager.
      Parameters:
      contextManager - The ContextManager to prune for.

  • Method Details

    • pruneMessages

      public void pruneMessages(List<Long> sequentialIds, String reason)
      Prunes entire ChatMessage objects from the context based on their unique sequential IDs.

      This method identifies all parts within the targeted messages and delegates the actual removal to prunePartsByReference(List, String) to ensure all dependencies are also pruned.

      Parameters:
      sequentialIds - A list of message sequential IDs to remove.
      reason - The reason for the pruning.

    • pruneParts

      public void pruneParts(long messageSequentialId, List<Number> partIndices, String reason)
      Prunes specific parts from a single ChatMessage, identified by the message's sequential ID and the parts' indices.

      This method identifies the specific Part objects and delegates the removal to prunePartsByReference(List, String) to handle dependency resolution.

      Parameters:
      messageSequentialId - The sequential ID of the message containing the parts to prune.
      partIndices - A list of zero-based indices of the parts to remove.
      reason - The reason for the pruning.

    • prunePartsByReference

      public void prunePartsByReference(List<com.google.genai.types.Part> initialCandidates, String reason)
      The low-level workhorse method that removes a given set of Part objects from the context.

      This method performs a recursive, cross-message dependency traversal to find all parts that must be removed along with the initial candidates. It then updates the affected messages and cleans their dependency maps.

      Parameters:
      initialCandidates - A list of the actual Part objects to remove.
      reason - The reason for the pruning.

    • pruneEphemeralToolCall

      public void pruneEphemeralToolCall(String toolCallId, String reason)
      Prunes a single ephemeral tool call and its associated response.
      Parameters:
      toolCallId - The unique ID of the tool call.
      reason - The reason for pruning.

    • pruneEphemeralToolCalls

      public void pruneEphemeralToolCalls(List<String> toolCallIds, String reason)
      Prunes a list of ephemeral tool calls and their associated responses.
      Parameters:
      toolCallIds - The list of tool call IDs.
      reason - The reason for pruning.
      Throws:
      IllegalArgumentException - if any of the IDs are associated with a stateful resource.

    • pruneOther

      public void pruneOther(List<com.google.genai.types.Part> parts, String reason)
      Prunes non-tool parts (text, blobs, etc.) from the context.
      Parameters:
      parts - The list of Part objects to prune.
      reason - The reason for pruning.
      Throws:
      IllegalArgumentException - if any of the parts are tool calls or responses.

    • pruneToolCall

      public void pruneToolCall(String toolCallId, String reason)
      Prunes a single tool call (ephemeral or stateful) by its ID.
      Parameters:
      toolCallId - The tool call ID.
      reason - The reason for pruning.

    • pruneToolCalls

      public void pruneToolCalls(List<String> toolCallIds, String reason)
      Prunes a list of tool calls by their IDs.
      Parameters:
      toolCallIds - The list of tool call IDs.
      reason - The reason for pruning.

    • pruneEphemeralToolCalls

      public void pruneEphemeralToolCalls(ToolManager toolManager)
      Implements the automatic pruning of old tool calls based on the "ephemeral" rules.

      A part is considered ephemeral if:

      • The tool is explicitly marked as ContextBehavior.EPHEMERAL.
      • It is a function call with no corresponding response (orphan).
      • It is a function response from a stateful tool that failed to return a resource.
      This method scans messages older than 5 user turns and prunes any ephemeral parts found.

      Parameters:
      toolManager - The ToolManager, needed to check tool metadata.