fix bug of chat without stream (infiniflow#830)

### What problem does this PR solve? ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
aigc-hf · May 17, 2024 · 673a28e · 673a28e
1 parent 2bfacd0
commit 673a28e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 13 deletions.
diff --git a/api/apps/api_app.py b/api/apps/api_app.py
@@ -222,10 +222,13 @@ def stream():
             resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
             return resp
         else:
-            ans = chat(dia, msg, **req)
-            fillin_conv(ans)
-            API4ConversationService.append_message(conv.id, conv.to_dict())
-            return get_json_result(data=ans)
+            answer = None
+            for ans in chat(dia, msg, **req):
+                answer = ans
+                fillin_conv(ans)
+                API4ConversationService.append_message(conv.id, conv.to_dict())
+                break
+            return get_json_result(data=answer)
 
     except Exception as e:
         return server_error_response(e)

diff --git a/api/apps/conversation_app.py b/api/apps/conversation_app.py
@@ -162,10 +162,13 @@ def stream():
             return resp
 
         else:
-            ans = chat(dia, msg, **req)
-            fillin_conv(ans)
-            ConversationService.update_by_id(conv.id, conv.to_dict())
-            return get_json_result(data=ans)
+            answer = None
+            for ans in chat(dia, msg, **req):
+                answer = ans
+                fillin_conv(ans)
+                ConversationService.update_by_id(conv.id, conv.to_dict())
+                break
+            return get_json_result(data=answer)
     except Exception as e:
         return server_error_response(e)
 
diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py
@@ -84,8 +84,7 @@ def chat(dialog, messages, stream=True, **kwargs):
     kbs = KnowledgebaseService.get_by_ids(dialog.kb_ids)
     embd_nms = list(set([kb.embd_id for kb in kbs]))
     if len(embd_nms) != 1:
-        if stream:
-            yield {"answer": "**ERROR**: Knowledge bases use different embedding models.", "reference": []}
+        yield {"answer": "**ERROR**: Knowledge bases use different embedding models.", "reference": []}
         return {"answer": "**ERROR**: Knowledge bases use different embedding models.", "reference": []}
 
     questions = [m["content"] for m in messages if m["role"] == "user"]
@@ -126,8 +125,7 @@ def chat(dialog, messages, stream=True, **kwargs):
         "{}->{}".format(" ".join(questions), "\n->".join(knowledges)))
 
     if not knowledges and prompt_config.get("empty_response"):
-        if stream:
-            yield {"answer": prompt_config["empty_response"], "reference": kbinfos}
+        yield {"answer": prompt_config["empty_response"], "reference": kbinfos}
         return {"answer": prompt_config["empty_response"], "reference": kbinfos}
 
     kwargs["knowledge"] = "\n".join(knowledges)
@@ -177,7 +175,7 @@ def decorate_answer(answer):
                 **kwargs), msg, gen_conf)
         chat_logger.info("User: {}|Assistant: {}".format(
             msg[-1]["content"], answer))
-        return decorate_answer(answer)
+        yield decorate_answer(answer)
 
 
 def use_sql(question, field_map, tenant_id, chat_mdl, quota=True):