Merge pull request #191 from NVIDIA/feature/fix-chat-generation-chunk…

…-issue Fix ChatGenerationChunk issue.
NVIDIA · Dec 20, 2023 · 7e0737c · 7e0737c
2 parents f19a85f + d4de775
commit 7e0737c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemoguardrails/streaming.py b/nemoguardrails/streaming.py
@@ -191,6 +191,8 @@ async def push_chunk(
             chunk = chunk.text
         elif isinstance(chunk, AIMessageChunk):
             chunk = chunk.content
+        elif isinstance(chunk, ChatGenerationChunk):
+            chunk = chunk.text
         elif isinstance(chunk, str) or chunk is None:
             pass
         else: