如何用拦截器模拟大模型返回的数据

前言

通过之前在 Android 处理流式响应 - 掘金 (juejin.cn) 中的介绍，我们已经了解了使用 OkHttp 时 对于流式响应的处理方式。OkHttp 的一大特点就是提供了非常灵活的拦截器机制，日常开发中可以通过设置拦截器返回特定的数据，实现接口响应 mock 的功能。那么对于流式响应，由于数据不是一次性返回的，拦截器又该如何实现呢？

OkHttp 的拦截器

对于传统的 http 接口响应，使用 OkHttp 的拦截器进行数据 mock 是很简单的。

class SkipNetworkInterceptor : Interceptor {
    
    val gson = Gson()
    
    override fun intercept(chain: Interceptor.Chain): Response {
        val mockResult = "mock data"
        return Response.Builder()
        .code(200)
        .request(chain.request())
        .protocol(Protocol.HTTP_1_1)
        .message("OK")
        .body(
            gson.toJson(mockResult).toResponseBody("application/json".toMediaType())
        ).build()
    }
}

实现 intercept 方法，返回一个 Response 即可。

这里的重点其实是 body 这个部分，对于非流式响应，返回结果无论多么复杂、多么大。总是可以构建出一个最终的Response，直接返回就好了。但是，对于流式响应，body 这个部分又该如何构建呢？对于一次请求来说，响应内容是不断变化的，但是 intercept 方法一旦返回就结束了，body 不断变化的部分该如何实现呢？

ResponseBody

我们先从上面的 SkipNetworkInterceptor 这个拦截器出发，看看常规的 Response 是如何创建的。

toResponseBody

以这里 Json 字符串的扩展函数 toResponseBody 为例

    fun String.toResponseBody(contentType: MediaType? = null): ResponseBody {
      var charset: Charset = UTF_8
      var finalContentType: MediaType? = contentType
      if (contentType != null) {
        val resolvedCharset = contentType.charset()
        if (resolvedCharset == null) {
          charset = UTF_8
          finalContentType = "$contentType; charset=utf-8".toMediaTypeOrNull()
        } else {
          charset = resolvedCharset
        }
      }
      val buffer = Buffer().writeString(this, charset)
      return buffer.asResponseBody(finalContentType, buffer.size)
    }

这个方法主要做了以下几件事情

• 根据 contentType 确定了字符编码及最终的 contentType
• 创建 一个 Buffer 对象，将当前字符串的内容基于字符编码写入到这个 buffer 中。
• 调用asResponseBody 扩展函数返回一个 Response

asResponseBody

    fun BufferedSource.asResponseBody(
      contentType: MediaType? = null,
      contentLength: Long = -1L
    ): ResponseBody = object : ResponseBody() {
      override fun contentType() = contentType

      override fun contentLength() = contentLength

      override fun source() = this@asResponseBody
    }

asResponseBody 的实现就比较简单了，基于当前 Buffer 对象的大小和 contentType 实现 ResponseBody 抽象方法即可，同时 source 方法返回了当前 Buffer 对象。

通过上面的实现可以看到，创建 ResponseBody 时需要定义原始数据的大小、contentType，并提供一个 BufferedSource 的实例，这个实例负责存储上游写入的数据，同时下游会基于这个实例读取写入的数据。

Pipe

理清楚了构建 ResponseBody 的要点，我们就可以实现基于流式响应的 ResponseBody 了，这里我们需要借助 Pipe 这个类。

/**
 * A source and a sink that are attached. The sink's output is the source's input. Typically each
 * is accessed by its own thread: a producer thread writes data to the sink and a consumer thread
 * reads data from the source.
 *
 */
class Pipe(internal val maxBufferSize: Long) { ...}

Pipe 顾名思义，就是管道，他的作用和 | 这个管道操作符的功能非常相似，Pipe 内部维护了 source 和 sink ，source 会将 sink 的输出当做输入，是一个典型的生产者消费者模式。

Pipe 这个类的核心功能可以抽象成下面这样

class Pipe(internal val maxBufferSize: Long) {
  internal val buffer = Buffer()

  @get:JvmName("sink")
  val sink = object : Sink {

    override fun write(source: Buffer, byteCount: Long) {
      lock.withLock {
        while (byteCount > 0) {
          val bytesToWrite = minOf(bufferSpaceAvailable, byteCount)
          buffer.write(source, bytesToWrite)
        }
      }
    }
  }

  @get:JvmName("source")
  val source = object : Source {
    override fun read(sink: Buffer, byteCount: Long): Long {
        val result = buffer.read(sink, byteCount)
        return result
      }
    }
  }
}

内部创建了一个 Buffer 的实例 buffer ，通过 sink 这个接口实现数据的写操作，source 实现数据的读取操作，读和写都是在操作 buffer 这个公共的缓冲池。当然实际代码还要处理读写同步、数据流的关闭等异常边界的 case。

对于我们关注的问题，看懂了这些，想要实现流式响应的 ResponseBody 就有谱了。

流式响应的Interceptor

创建 ResponseBody

    val pipe = Pipe(8192)


    private val responseBody = object : ResponseBody() {
        override fun contentLength(): Long {
            return -1
        }

        override fun contentType(): MediaType? {
            return "text/event-stream; charset=utf-8".toMediaTypeOrNull()
        }

        override fun source(): BufferedSource {
            return pipe.source.buffer()
        }
    }

这里我们创建了一个 Pipe 的实例。contentType 按照 SSE 的标准定义进行约束，source 方法直接返回 pipe 中 source.buffer() 即可。需要注意的是这里 source.buffer 通过扩展函数返回的其实是 RealBufferedSource 。

fun Source.buffer(): BufferedSource = RealBufferedSource(this)

流式数据的写入

按照 Pipe 的定义，buffer 中缓存的数据会由发起请求的客户端在接收到 OkHttp 的 Response 时，通过 source 接口主动读取，我们需要通过 sink 接口将数据写入到 buffer 中。严格来说是 RealBufferedSink ,实际写入时并不是直接使用 sink 的 write 方法，而是基于其封装过后的 RealBufferedSink 进行写操作。这里 RealBufferedSink 和上面的 RealBufferedSource 通过 Buffer 这层包装类，提升了读写操作的效率，避免了频繁读写小数据的问题，这也是 IO 操作的一个共识。

    val content =
        "在Java中，Timer 类是java.util.Timer的简称，它是一个用于安排任务以后在后台线程中执行的工具。Timer可以安排一个java.util.TimerTask任务，以固定延迟或者固定周期重复执行"
    val contentArray = content.toCharArray()
    var index = 0
    val moshi: Moshi = Moshi.Builder().add(KotlinJsonAdapterFactory()).build()
    val jsonAdapter: JsonAdapter<SSResponse> = moshi.adapter(SSResponse::class.java)

    val sink = pipe.sink.buffer()

    private val task = object : TimerTask() {
        override fun run() {
            val input = contentArray[index]
            val end = index >= contentArray.size - 1
            val sseResult = SSEResult(System.currentTimeMillis().toString(), input.toString())
            val ssrResponse = SSResponse(200, "OK", end, sseResult)
            val json: String = jsonAdapter.toJson(ssrResponse)
            val mock = "data: $json\n\n"

            sink.writeUtf8(mock)
            sink.flush()
            index++
            if (index >= contentArray.size) {
                sink.close()
            }
        }
    }

这里我们定义一个写数据的任务，简单起见就是通过遍历字符串，每次写入一个字符，模拟流式响应的效果，这里需要关注以下几点。

• 数据格式：val mock = "data: $json\n\n" 这里要和请求侧解析数据的格式保持匹配，最好是遵循 SSE 的协议。
• 调用 sink.flush 方法，实现每一包数据的立刻写入。
• 数据写操作完成后一定要关闭 sink。

intercept

    val timer: Timer = Timer()

    override fun intercept(chain: Interceptor.Chain): Response {
        val request = chain.request()

        if (url.contains("stream_chat")) {

            timer.schedule(task, 0, 20)

            return Response.Builder().request(request).code(200).message("OK")
                .protocol(okhttp3.Protocol.HTTP_1_1).addHeader("Content-Type", "text/event-stream")
                .body(responseBody).build()
        }
        return chain.proceed(request)
    }

这里我们通过 timer 实现对上述任务每 20 毫秒的一次调用，在 intercept 方法中返回之前的定义好的 responseBody 即可。

响应读取

我们在 OkHttp Client 的定义中添加这个拦截器之后，可以看看一下输出

private fun sseCall() {
    val sb = StringBuilder()
    OkHttpUtil.sseHandler<SSResponse> { it, k ->
        if (it != null) {
            sb.append(it.event.data)
        }
        log(sb.toString())
        if (Objects.isNull(k).not()) {
            log(k!!)
        }
    }
}

OkHttpUtil:onOpen() called with: eventSource = okhttp3.internal.sse.RealEventSource@52a56147, response = Response{protocol=http/1.1, code=200, message=OK, url=http://localhost:8199/stream_chat}
在
在J
在Ja
在Jav
在Java
在Java中
在Java中，
在Java中，T
在Java中，Ti
在Java中，Tim
在Java中，Time
在Java中，Timer
在Java中，Timer 
在Java中，Timer 类
在Java中，Timer 类是
在Java中，Timer 类是j
在Java中，Timer 类是ja
在Java中，Timer 类是jav
在Java中，Timer 类是java
在Java中，Timer 类是java.
...
在Java中，Timer 类是java.util.Timer的简称，它是一个用于安排任务以后在后台线程中执行的工具。Timer可以安排一个java.util.TimerTask任务，以固定延迟或者固定周期重复执行
closed

可以看到最终结果已经按照单个字符依次返回了。

-1

最后，我们再来看看这里实现 responseBody 时的一个细节 contentLength() 为什么返回 -1 。对于流式响应来说，数据的大小在响应首次返回的那一刻是无法确定的，那么对于这个 -1 在 OkHttp 内部又是如何处理的呢？

我们直接看 OkHttp 拦截器集合中的最后一个 CallServerInterceptor, 看看他最终是如何处理上一个响应的。最终的 Response 会由 Exchange.openResponseBody 处理。

  override fun openResponseBodySource(response: Response): Source {
    return when {
      !response.promisesBody() -> newFixedLengthSource(0)
      response.isChunked -> newChunkedSource(response.request.url)
      else -> {
        val contentLength = response.headersContentLength()
        if (contentLength != -1L) {
          newFixedLengthSource(contentLength)
        } else {
          newUnknownLengthSource()
        }
      }
    }
  }

可以看到对于 contentLength 是否等于 -1 有两种不同的处理情况。我们先看不等于 -1 的场景

contentLenght!=-1

  private inner class FixedLengthSource(private var bytesRemaining: Long) :
      AbstractSource() {

    init {
      if (bytesRemaining == 0L) {
        responseBodyComplete()
      }
    }

    override fun read(sink: Buffer, byteCount: Long): Long {
      require(byteCount >= 0L) { "byteCount < 0: $byteCount" }
      check(!closed) { "closed" }
      if (bytesRemaining == 0L) return -1

      val read = super.read(sink, minOf(bytesRemaining, byteCount))
      if (read == -1L) {
        connection.noNewExchanges() // The server didn't supply the promised content length.
        val e = ProtocolException("unexpected end of stream")
        responseBodyComplete()
        throw e
      }

      bytesRemaining -= read
      if (bytesRemaining == 0L) {
        responseBodyComplete()
      }
      return read
    }

    override fun close() {
      if (closed) return

      if (bytesRemaining != 0L &&
          !discard(ExchangeCodec.DISCARD_STREAM_TIMEOUT_MILLIS, MILLISECONDS)) {
        connection.noNewExchanges() // Unread bytes remain on the stream.
        responseBodyComplete()
      }

      closed = true
    }
  }

contentLenght == -1

  private inner class UnknownLengthSource : AbstractSource() {
    private var inputExhausted: Boolean = false

    override fun read(sink: Buffer, byteCount: Long): Long {
      require(byteCount >= 0L) { "byteCount < 0: $byteCount" }
      check(!closed) { "closed" }
      if (inputExhausted) return -1

      val read = super.read(sink, byteCount)
      if (read == -1L) {
        inputExhausted = true
        responseBodyComplete()
        return -1
      }
      return read
    }

    override fun close() {
      if (closed) return
      if (!inputExhausted) {
        responseBodyComplete()
      }
      closed = true
    }
  }

可以看到这两种情况最核心的差异还是对于请求是否完成的处理逻辑不一样，对于有固定长度的响应，按照读取的内容做从总长度中每次减掉一部分，直到最后为 0 即可。而对于 contentLength = -1 的这种长度未知的情况，只能是基于读取的内容是否到结尾来判断。

具体代码可以参考 Github

小结

通过实现流式响应拦截器的功能，对 OkHttp Response 的实现链路又有了更多的认知。再次感受到 OkHttp 的代码写的真的太巧妙了，很多小细节的处理真的很值得借鉴。