refactor(SSEProcessor): return early on '[DONE]' chunk detection

unit-mesh · Oct 23, 2024 · e801132 · e801132
1 parent 3112821
commit e801132
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomSSEProcessor.kt b/core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomSSEProcessor.kt
@@ -92,6 +92,10 @@ open class CustomSSEProcessor(private val project: Project) {
                                 // {"id":"cmpl-a22a0d78fcf845be98660628fe5d995b","object":"chat.completion.chunk","created":822330,"model":"moonshot-v1-8k","choices":[{"index":0,"delta":{},"finish_reason":"stop","usage":{"prompt_tokens":434,"completion_tokens":68,"total_tokens":502}}]}
                                 // in some case, the response maybe not equal to our response format, so we need to ignore it
                                 // {"id":"cmpl-ac26a17e","object":"chat.completion.chunk","created":1858403,"model":"yi-34b-chat","choices":[{"delta":{"role":"assistant"},"index":0}],"content":"","lastOne":false}
+                                if (sse.data == "[DONE]") {
+                                    return@blockingForEach
+                                }
+
                                 val chunk: String? = JsonPath.parse(sse!!.data)?.read(responseFormat)
 
                                 // new JsonPath lib caught the exception, so we need to handle when it is null