Spring AI Alibaba 24｜Hooks 与 Interceptors · 自定义 Hooks 和 Interceptors：ModelHook

import com.alibaba.cloud.ai.graph.agent.hook.ModelHook;
import com.alibaba.cloud.ai.graph.agent.hook.HookPosition;
import com.alibaba.cloud.ai.graph.agent.hook.HookPositions;
import java.util.concurrent.CompletableFuture;

@HookPositions({HookPosition.BEFORE_MODEL, HookPosition.AFTER_MODEL})
public class CustomModelHook extends ModelHook {

  @Override
  public String getName() {
      return "custom_model_hook";
  }

  @Override
  public CompletableFuture<Map<String, Object>> beforeModel(OverAllState state, RunnableConfig config) {
      // 在模型调用前执行
      System.out.println("准备调用模型...");

      // 可以修改状态
      // 例如：添加额外的上下文
      return CompletableFuture.completedFuture(Map.of("extra_context", "某些额外信息"));
  }

  @Override
  public CompletableFuture<Map<String, Object>> afterModel(OverAllState state, RunnableConfig config) {
      // 在模型调用后执行
      System.out.println("模型调用完成");

      // 可以记录响应信息
      return CompletableFuture.completedFuture(Map.of());
  }
}

import com.alibaba.cloud.ai.graph.agent.hook.ModelHook;
import com.alibaba.cloud.ai.graph.agent.hook.HookPosition;
import com.alibaba.cloud.ai.graph.agent.hook.HookPositions;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.state.RemoveByHash;
import org.springframework.ai.chat.messages.Message;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;

@HookPositions({HookPosition.BEFORE_MODEL})
public class MessageDeletionHook extends ModelHook {

  @Override
  public CompletableFuture<Map<String, Object>> beforeModel(OverAllState state, RunnableConfig config) {
      Optional<Object> messagesOpt = state.value("messages");
      if (!messagesOpt.isPresent()) {
          return CompletableFuture.completedFuture(Map.of());
      }

      List<Message> messages = (List<Message>) messagesOpt.get();
      
      // 构建新的消息列表，保持原顺序
      List<Object> newMessages = new ArrayList<>();
      for (Message msg : messages) {
          // 根据条件决定保留或删除
          if (shouldKeep(msg)) {
              newMessages.add(msg);  // 保留消息
          } else {
              newMessages.add(RemoveByHash.of(msg));  // 标记删除
          }
      }
      
      return CompletableFuture.completedFuture(Map.of("messages", newMessages));
  }
  
  private boolean shouldKeep(Message msg) {
      // 实现你的保留逻辑
      return true;
  }
}