fix streaming chunk

InternLM · Oct 24, 2024 · bc4ebf6 · bc4ebf6
1 parent 573c7b5
commit bc4ebf6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lagent/llms/openai.py b/lagent/llms/openai.py
@@ -267,7 +267,7 @@ def streaming(raw_response):
                     chunk_size=8192, decode_unicode=False, delimiter=b'\n'):
                 if chunk:
                     decoded = chunk.decode('utf-8')
-                    if decoded == 'data: [DONE]':
+                    if decoded.startswith('data: [DONE]'):
                         return
                     if decoded[:5] == 'data:':
                         decoded = decoded[5:]
@@ -716,7 +716,7 @@ async def streaming(raw_response):
             async for chunk in raw_response.content:
                 if chunk:
                     decoded = chunk.decode('utf-8')
-                    if decoded == 'data: [DONE]':
+                    if decoded.startswith('data: [DONE]'):
                         return
                     if decoded[:5] == 'data:':
                         decoded = decoded[5:]