Fix Doc-Sum stream output format

Workround to keep Doc-Sum stream output aligned with v1.1 format This is a workaround to extract the tokens from stream output. Fix issue: opea-project/GenAIInfra#753 Signed-off-by: Wang, Xigui <[email protected]>
opea-project · Jan 24, 2025 · 4f45127 · 4f45127
1 parent c4362a4
commit 4f45127
Showing 1 changed file with 19 additions and 5 deletions.
diff --git a/comps/llms/src/doc-summarization/integrations/common.py b/comps/llms/src/doc-summarization/integrations/common.py
@@ -190,15 +190,29 @@ async def generate(self, input: DocSumChatCompletionRequest, client):
 
         if input.stream:
 
-            async def stream_generator():
-                from langserve.serialization import WellKnownLCSerializer
+            import json
+            from langserve.serialization import WellKnownLCSerializer
+            _serializer = WellKnownLCSerializer()
+            def extract_llm_tokens(stream_output):
+                op_data = _serializer.dumps({"ops": stream_output.ops}).decode("utf-8")
+                logger.info(op_data)
+                parsed_data = json.loads(op_data)
+
+                tokens = []
+                for op in parsed_data.get("ops", []):
+                    if op["op"] == "add" and "/streamed_output_str" in op["path"]:
+                        tokens.append(op["value"])
+
+                return "".join(tokens)
 
-                _serializer = WellKnownLCSerializer()
+            async def stream_generator():
                 async for chunk in llm_chain.astream_log(docs):
-                    data = _serializer.dumps({"ops": chunk.ops}).decode("utf-8")
+                    data = extract_llm_tokens(chunk)
+
                     if logflag:
                         logger.info(data)
-                    yield f"data: {data}\n\n"
+                    if data != '':
+                        yield f"data: {data}\n\n"
                 yield "data: [DONE]\n\n"
 
             return StreamingResponse(stream_generator(), media_type="text/event-stream")