流式输出——解决 HITL 难题 (SpringAIAlibaba)

前言

在基于 SpringAIAlibaba 的应用开发中，有一个潜力很大的高级功能 HITL(人工介入)，它允许用户中断Agent的运行，在做出一些行为之后，再恢复Agent的运行任务。可以修改工具调用，可以输入补充信息来调整Agent的运行方向.....
但是当我们将它中断后长时间没有给它回应 时，在流式输出下，可能会提前结束这次响应 ，导致效果不对。

这篇文章，以 HITL 中 工具调用产生的中断为例子，讲述一下如何解决。

快速结论

使用 Sinks.Many<ReturnValue> 类的流式数据推送替代 Flux<ReturnValue> 的返回，在整套流程结束后再关闭。
关键点 ：fluxInstance.subscribe() ， sinksManyInstance.tryEmitNext()

详细流程

正常情况：

一般情况下，代码可能表现为这样：

@GetMapping(value = "/stream/chat",
       produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public Flux<String> testStreamHumanInLoop() throws GraphRunnerException {

   String threadId = "stream-session-123";

   Flux<NodeOutput> outputFlux = dashscopeHITLAgent.stream("考虑调用工具，帮我看看可乐鸡翅的配方",
           RunnableConfig.builder().threadId(threadId).build());

  return convertToStringFlux(outputFlux);

}   

private Flux<String>  convertToStringFlux(Flux<NodeOutput> outputFlux){
   return outputFlux.flatMap(out -> {
if (out instanceof StreamingOutput<?> streamingOutput) {
           // 处理流式输出
           String chunk = streamingOutput.chunk();
           if (chunk != null && !chunk.isEmpty()) {
               SystemPrinter.println("流式chunk: " + chunk);
               return Flux.just(chunk);
           }
       }
       return Flux.empty();
   });
}

这种流式输出，在大部分的无中断场景 中都可以正常运行，但是出现了中断后，flux 会关闭推送，导致请求结束。
目前有两种 解决方式：

方案一

请求结束后，工具调用进入待审核状态，保存相关信息，在审核结束后，恢复Agent运行，并返回剩余的 Flux 对象。即：（示例中用 volatile模拟，实际需要 将 interruptionMetaData 和 threadId 绑定到数据库等）

public volatile InterruptionMetadata testHitldata;

private Flux<String>  convertToStringFlux(Flux<NodeOutput> outputFlux){
 return outputFlux.flatMap(out -> {
     if (out instanceof InterruptionMetadata interruptionMetadata) {
         // 处理人工介入中断
         // 临时存储 
         testHitldata = interruptionMetadata;
         return Flux.just("[系统] 发生了工具调用中断，请等待人工审核结果...");
     } else if (out instanceof StreamingOutput<?> streamingOutput) {
         // 处理流式输出
         String chunk = streamingOutput.chunk();
         if (chunk != null && !chunk.isEmpty()) {
             SystemPrinter.println("流式chunk: " + chunk);
             return Flux.just(chunk);
         }
     }
     return Flux.empty();
 });
}



@GetMapping(value = "/approve", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public Flux<String> approve() throws GraphRunnerException {

 
 String threadId = "stream-session-123";

 InterruptionMetadata approvalMetadata = getReviewResult(testHitldata);

 RunnableConfig resumeConfig = RunnableConfig.builder()
         .threadId(threadId)
         .addMetadata(RunnableConfig.HUMAN_FEEDBACK_METADATA_KEY, approvalMetadata)
         .build();
         
 Flux<NodeOutput> resumeFlux = dashscopeHITLAgent.stream(" ", resumeConfig);
 // 重置 testHitldata
 testHitldata = null;

 return convertToStringFlux(resumeFlux);

}

拥有一定开发经验的老家伙们 肯定看出来了，这种方法 仅适用于 单个用户 审查的场景。比如 AI IDE中试图执行命令时需要经过用户同意才能继续 (实现的底层逻辑不一样，但场景很相似)

当中断可能由不同用户审核 时，这个就没用了。 所以看 方案二 吧。

方案二：

使用 Sinks.Many 类进行消息的推送，Flux 仅仅作为生产方 ，朝 sinkManyInstance发送消息 ，由sink进行传递。如以下示例代码:

@GetMapping(value = "/stream/humanInLoop",
      produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public Flux<String> testStreamHumanInLoop() throws GraphRunnerException {

  String threadId = "stream-session-123";

  // 创建一个 sink 发送消息
  Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();
  memoryCacheService.cacheInterruptSink(threadId, sink);

  Flux<NodeOutput> outputFlux = dashscopeHITLAgent.stream("考虑调用工具，帮我看看可乐鸡翅的配方",
          RunnableConfig.builder().threadId(threadId).build());

  outputFlux.subscribe(
          out -> handleFluxChunk(threadId, sink, out),
          sink::tryEmitError,
          () -> sinkOnComplete(threadId, sink));

  return sink.asFlux();

}

private void handleFluxChunk(String threadId, Sinks.Many<String> sink, NodeOutput out) {
  if (out instanceof InterruptionMetadata interruptionMetadata) {
      for (InterruptionMetadata.ToolFeedback toolFeedback : interruptionMetadata.toolFeedbacks()) {
          String toolName = toolFeedback.getName();

          SystemPrinter.println("工具调用中断: " + toolName);
      }
      if (memoryCacheService.hasInterruptSink(threadId)) {
          sink.tryEmitNext("[系统] 发生了工具调用中断，请等待人工审核结果...");
          testHitldata = interruptionMetadata;
      } else {
          sink.tryEmitNext("[系统] 系统出现异常，无法处理人工审核，请稍后再试...");
          sink.tryEmitComplete();
      }
  }
  if (out instanceof StreamingOutput<?> streamingOutput) {

      String chunk = streamingOutput.chunk();

      if (chunk != null && !chunk.isEmpty()) {

          SystemPrinter.println("流式chunk: " + chunk);

          sink.tryEmitNext(chunk);
      }
  }
}


private void sinkOnComplete(String threadId, Sinks.Many<String> sink) {
  if (testHitldata == null) {
      SystemPrinter.println("任务真正完成，关闭流");
      sink.tryEmitComplete();
      memoryCacheService.removeInterruptSink(threadId);
  } else {
      SystemPrinter.println("检测到中断，流保持开启，等待审核...");
  }
}

上面的示例代码 实际上就是 我认为的 最优解了，实际场景中 interruptionMetaData和 threadId 都是会被存储到数据库 中的，审核的时候提取出来 就能完成整套流程了。

sink 关键解释

细心的朋友们肯定看出来了，在示例代码中，sink 被我存放到了 memoryCacheService(内存缓存服务) 中，那为什么不能缓存到最常用的redis中呢？
我们先看看 redis中它的效果是什么：

{"cancelled":false,"scanAvailable":true}

怎么只有这么一小点... 比我随便写的一个 DTO 都要小了。
这是因为 sink 对象中有大量的 volatile 变量，它实质上是存在于JVM 中的一种管道设施 ，需要和JVM共同存在 ，存放到 redis中的 只是一个外壳 。自然从Redis 获取的时候就会报错 ，不能正确获取 。

所以必须存放到内存中，又由于 它需要线程隔离且和线程绑定 (Agent的运行线程)，所以 存放它的地方应该是ConcurrentHashMap ,这里给出我的代码示例:

@Service
public class MemoryCacheService {

    private final Map<String, Sinks.Many<String>> threadInterruptSinks = new ConcurrentHashMap<>();

    public void cacheInterruptSink(String threadId, Sinks.Many<String> sink) {
        threadInterruptSinks.put(threadId, sink);
    }

    public boolean hasInterruptSink(String threadId) {
        return threadInterruptSinks.containsKey(threadId);
    }
    public Sinks.Many<String> getInterruptSink(String threadId) {
        return threadInterruptSinks.get(threadId);
    }



    public void removeInterruptSink(String threadId) {
        threadInterruptSinks.remove(threadId);
    }
}

方案二 就是真的解决了 方案一不能解决的 不同用户问题 ，因为Sinks.Many支持多端接收数据 ，即不同的线程也可以往里面传输数据(volatile变量)，在发送数据 的时候，它只是发送给了 前端(后端返回值sink.asFlux() 也是一个 Flux 对象).

其他问题

返回值既然也是flux对象，那为什么在遇到中断事件 的时候，没有结束呢？
让我们看看 示例代码中的 onSinkComplete() 方法:

private void sinkOnComplete(String threadId, Sinks.Many<String> sink) {
    if (testHitldata == null) {
        SystemPrinter.println("任务真正完成，关闭流");
        sink.tryEmitComplete();
        memoryCacheService.removeInterruptSink(threadId);
    } else {
        SystemPrinter.println("检测到中断，流保持开启，等待审核...");
    }
}

它是在 testHitldata == null 的时候 才允许结束，这个方法就是在中断的时候 flux 会调用的，此时flux认为已经没有数据了，就会尝试调用 onComplete()，也就是我们填入的 这个方法。

后续疑问

那在方案一中也这样 可不可以防止flux关闭呢？ 自己去试试吧，博主这里就不揭晓谜底了。
tip: 问题的关键不在这里

总结

1. 对于 HITL 在流式输出上出现的解决方法，有两种。
2. 第一种适合单个用户的情况。
3. 第二种适合多个用户的情况。

结尾疑问

对于第二种方案，如果是在分布式环境 下，应该如何解决恢复请求被发送到另外一个JVM上了 ，导致这个JVM的sink 无法收到消息的问题呢？评论区给出你的答案叭，嘻嘻。