add stream chat (infiniflow#811)

### What problem does this PR solve?

infiniflow#709 
### Type of change

- [x] New Feature (non-breaking change which adds functionality)

api/apps/api_app.py

@@ -13,10 +13,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import json
 import os
 import re
 from datetime import datetime, timedelta
-from flask import request
+from flask import request, Response
 from flask_login import login_required, current_user
 
 from api.db import FileType, ParserType
@@ -31,11 +32,11 @@
 from api.utils import get_uuid, current_timestamp, datetime_format
 from api.utils.api_utils import server_error_response, get_data_error_result, get_json_result, validate_request
 from itsdangerous import URLSafeTimedSerializer
-from api.db.services.task_service import TaskService, queue_tasks
+
 from api.utils.file_utils import filename_type, thumbnail
 from rag.utils.minio_conn import MINIO
-from api.db.db_models import Task
-from api.db.services.file2document_service import File2DocumentService
+
+
 def generate_confirmation_token(tenent_id):
     serializer = URLSafeTimedSerializer(tenent_id)
     return "ragflow-" + serializer.dumps(get_uuid(), salt=tenent_id)[2:34]
@@ -164,6 +165,7 @@ def completion():
     e, conv = API4ConversationService.get_by_id(req["conversation_id"])
     if not e:
         return get_data_error_result(retmsg="Conversation not found!")
+    if "quote" not in req: req["quote"] = False
 
     msg = []
     for m in req["messages"]:
@@ -180,13 +182,45 @@ def completion():
             return get_data_error_result(retmsg="Dialog not found!")
         del req["conversation_id"]
         del req["messages"]
-        ans = chat(dia, msg, **req)
+
         if not conv.reference:
             conv.reference = []
-        conv.reference.append(ans["reference"])
-        conv.message.append({"role": "assistant", "content": ans["answer"]})
-        API4ConversationService.append_message(conv.id, conv.to_dict())
-        return get_json_result(data=ans)
+        conv.message.append({"role": "assistant", "content": ""})
+        conv.reference.append({"chunks": [], "doc_aggs": []})
+
+        def fillin_conv(ans):
+            nonlocal conv
+            if not conv.reference:
+                conv.reference.append(ans["reference"])
+            else: conv.reference[-1] = ans["reference"]
+            conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
+
+        def stream():
+            nonlocal dia, msg, req, conv
+            try:
+                for ans in chat(dia, msg, True, **req):
+                    fillin_conv(ans)
+                    yield "data:"+json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
+                API4ConversationService.append_message(conv.id, conv.to_dict())
+            except Exception as e:
+                yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
+                                            "data": {"answer": "**ERROR**: "+str(e), "reference": []}},
+                                           ensure_ascii=False) + "\n\n"
+            yield "data:"+json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
+
+        if req.get("stream", True):
+            resp = Response(stream(), mimetype="text/event-stream")
+            resp.headers.add_header("Cache-control", "no-cache")
+            resp.headers.add_header("Connection", "keep-alive")
+            resp.headers.add_header("X-Accel-Buffering", "no")
+            resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
+            return resp
+        else:
+            ans = chat(dia, msg, False, **req)
+            fillin_conv(ans)
+            API4ConversationService.append_message(conv.id, conv.to_dict())
+            return get_json_result(data=ans)
+
     except Exception as e:
         return server_error_response(e)
 
@@ -229,7 +263,6 @@ def upload():
         return get_json_result(
             data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
 
-
     file = request.files['file']
     if file.filename == '':
         return get_json_result(
@@ -253,7 +286,6 @@ def upload():
             location += "_"
         blob = request.files['file'].read()
         MINIO.put(kb_id, location, blob)
-
         doc = {
             "id": get_uuid(),
             "kb_id": kb.id,
@@ -266,42 +298,11 @@ def upload():
             "size": len(blob),
             "thumbnail": thumbnail(filename, blob)
         }
-
-        form_data=request.form
-        if "parser_id" in form_data.keys():
-            if request.form.get("parser_id").strip() in list(vars(ParserType).values())[1:-3]:
-                doc["parser_id"] = request.form.get("parser_id").strip()
         if doc["type"] == FileType.VISUAL:
             doc["parser_id"] = ParserType.PICTURE.value
         if re.search(r"\.(ppt|pptx|pages)$", filename):
             doc["parser_id"] = ParserType.PRESENTATION.value
-
-        doc_result = DocumentService.insert(doc)
-
+        doc = DocumentService.insert(doc)
+        return get_json_result(data=doc.to_json())
     except Exception as e:
         return server_error_response(e)
-
-    if "run" in form_data.keys():
-        if request.form.get("run").strip() == "1":
-            try:
-                info = {"run": 1, "progress": 0}
-                info["progress_msg"] = ""
-                info["chunk_num"] = 0
-                info["token_num"] = 0
-                DocumentService.update_by_id(doc["id"], info)
-                # if str(req["run"]) == TaskStatus.CANCEL.value:
-                tenant_id = DocumentService.get_tenant_id(doc["id"])
-                if not tenant_id:
-                    return get_data_error_result(retmsg="Tenant not found!")
-
-                #e, doc = DocumentService.get_by_id(doc["id"])
-                TaskService.filter_delete([Task.doc_id == doc["id"]])
-                e, doc = DocumentService.get_by_id(doc["id"])
-                doc = doc.to_dict()
-                doc["tenant_id"] = tenant_id
-                bucket, name = File2DocumentService.get_minio_address(doc_id=doc["id"])
-                queue_tasks(doc, bucket, name)
-            except Exception as e:
-                 return server_error_response(e)
-
-    return get_json_result(data=doc_result.to_json())

api/apps/conversation_app.py

@@ -13,12 +13,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-from flask import request
+from flask import request, Response, jsonify
 from flask_login import login_required
 from api.db.services.dialog_service import DialogService, ConversationService, chat
 from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 from api.utils import get_uuid
 from api.utils.api_utils import get_json_result
+import json
 
 
 @manager.route('/set', methods=['POST'])
@@ -103,9 +104,12 @@ def list_convsersation():
 
 @manager.route('/completion', methods=['POST'])
 @login_required
-@validate_request("conversation_id", "messages")
+#@validate_request("conversation_id", "messages")
 def completion():
     req = request.json
+    #req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
+    #    {"role": "user", "content": "上海有吗？"}
+    #]}
     msg = []
     for m in req["messages"]:
         if m["role"] == "system":
@@ -123,13 +127,45 @@ def completion():
             return get_data_error_result(retmsg="Dialog not found!")
         del req["conversation_id"]
         del req["messages"]
-        ans = chat(dia, msg, **req)
+
         if not conv.reference:
             conv.reference = []
-        conv.reference.append(ans["reference"])
-        conv.message.append({"role": "assistant", "content": ans["answer"]})
-        ConversationService.update_by_id(conv.id, conv.to_dict())
-        return get_json_result(data=ans)
+        conv.message.append({"role": "assistant", "content": ""})
+        conv.reference.append({"chunks": [], "doc_aggs": []})
+
+        def fillin_conv(ans):
+            nonlocal conv
+            if not conv.reference:
+                conv.reference.append(ans["reference"])
+            else: conv.reference[-1] = ans["reference"]
+            conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
+
+        def stream():
+            nonlocal dia, msg, req, conv
+            try:
+                for ans in chat(dia, msg, True, **req):
+                    fillin_conv(ans)
+                    yield "data:"+json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
+                ConversationService.update_by_id(conv.id, conv.to_dict())
+            except Exception as e:
+                yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
+                                            "data": {"answer": "**ERROR**: "+str(e), "reference": []}},
+                                           ensure_ascii=False) + "\n\n"
+            yield "data:"+json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
+
+        if req.get("stream", True):
+            resp = Response(stream(), mimetype="text/event-stream")
+            resp.headers.add_header("Cache-control", "no-cache")
+            resp.headers.add_header("Connection", "keep-alive")
+            resp.headers.add_header("X-Accel-Buffering", "no")
+            resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
+            return resp
+
+        else:
+            ans = chat(dia, msg, False, **req)
+            fillin_conv(ans)
+            ConversationService.update_by_id(conv.id, conv.to_dict())
+            return get_json_result(data=ans)
     except Exception as e:
         return server_error_response(e)
 

api/apps/system_app.py

@@ -0,0 +1,67 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License
+#
+from flask_login import login_required
+
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.utils.api_utils import get_json_result
+from api.versions import get_rag_version
+from rag.settings import SVR_QUEUE_NAME
+from rag.utils.es_conn import ELASTICSEARCH
+from rag.utils.minio_conn import MINIO
+from timeit import default_timer as timer
+
+from rag.utils.redis_conn import REDIS_CONN
+
+
+@manager.route('/version', methods=['GET'])
+@login_required
+def version():
+    return get_json_result(data=get_rag_version())
+
+
+@manager.route('/status', methods=['GET'])
+@login_required
+def status():
+    res = {}
+    st = timer()
+    try:
+        res["es"] = ELASTICSEARCH.health()
+        res["es"]["elapsed"] = "{:.1f}".format((timer() - st)*1000.)
+    except Exception as e:
+        res["es"] = {"status": "red", "elapsed": "{:.1f}".format((timer() - st)*1000.), "error": str(e)}
+
+    st = timer()
+    try:
+        MINIO.health()
+        res["minio"] = {"status": "green", "elapsed": "{:.1f}".format((timer() - st)*1000.)}
+    except Exception as e:
+        res["minio"] = {"status": "red", "elapsed": "{:.1f}".format((timer() - st)*1000.), "error": str(e)}
+
+    st = timer()
+    try:
+        KnowledgebaseService.get_by_id("x")
+        res["mysql"] = {"status": "green", "elapsed": "{:.1f}".format((timer() - st)*1000.)}
+    except Exception as e:
+        res["mysql"] = {"status": "red", "elapsed": "{:.1f}".format((timer() - st)*1000.), "error": str(e)}
+
+    st = timer()
+    try:
+        qinfo = REDIS_CONN.health(SVR_QUEUE_NAME)
+        res["redis"] = {"status": "green", "elapsed": "{:.1f}".format((timer() - st)*1000.), "pending": qinfo["pending"]}
+    except Exception as e:
+        res["redis"] = {"status": "red", "elapsed": "{:.1f}".format((timer() - st)*1000.), "error": str(e)}
+
+    return get_json_result(data=res)