issue #1153: add start chunk

smclab · Dec 13, 2024 · 7898b2b · 7898b2b
1 parent 504d76e
commit 7898b2b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/ai-packages/rag-module/app/rag/chain.py b/ai-packages/rag-module/app/rag/chain.py
@@ -105,6 +105,8 @@ def get_chain(
         rephrase_chain = rephrase_prompt | llm | parser
         question = rephrase_chain.invoke({"question": question})
 
+    yield json.dumps({"chunk": "", "type": "START"})
+
     for chunk in chain.stream({"question": question, "context": documents}):
         yield json.dumps({"chunk": chunk, "type": "CHUNK"})
 
@@ -320,6 +322,8 @@ def get_chat_chain(
     documents_id = set()
     citations = []
 
+    yield json.dumps({"chunk": "", "type": "START"})
+
     for chunk in result:
         if "answer" in chunk.keys():
             result_answer += chunk