Refactor API for document and session (infiniflow#2819)

### What problem does this PR solve?

Refactor API for document and session.

### Type of change


- [x] Refactoring

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>

Author: Feiue

api/apps/sdk/doc.py

@@ -20,81 +20,77 @@
 
 from api.db import StatusEnum
 from api.db.services.dialog_service import DialogService, ConversationService, chat
-from api.settings import RetCode
 from api.utils import get_uuid
-from api.utils.api_utils import get_data_error_result
-from api.utils.api_utils import get_json_result, token_required
+from api.utils.api_utils import get_error_data_result
+from api.utils.api_utils import get_result, token_required
 
-
-@manager.route('/save', methods=['POST'])
+@manager.route('/chat/<chat_id>/session', methods=['POST'])
 @token_required
-def set_conversation(tenant_id):
+def create(tenant_id,chat_id):
     req = request.json
-    conv_id = req.get("id")
-    if "assistant_id" in req:
-        req["dialog_id"] = req.pop("assistant_id")
-    if "id" in req:
-        del req["id"]
-        conv = ConversationService.query(id=conv_id)
-        if not conv:
-            return get_data_error_result(retmsg="Session does not exist")
-        if not DialogService.query(id=conv[0].dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
-            return get_data_error_result(retmsg="You do not own the session")
-        if req.get("dialog_id"):
-            dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
-            if not dia:
-                return get_data_error_result(retmsg="You do not own the assistant")
-        if "dialog_id" in req and not req.get("dialog_id"):
-            return get_data_error_result(retmsg="assistant_id can not be empty.")
-        if "message" in req:
-            return get_data_error_result(retmsg="message can not be change")
-        if "reference" in req:
-            return get_data_error_result(retmsg="reference can not be change")
-        if "name" in req and not req.get("name"):
-            return get_data_error_result(retmsg="name can not be empty.")
-        if not ConversationService.update_by_id(conv_id, req):
-            return get_data_error_result(retmsg="Session updates error")
-        return get_json_result(data=True)
-
-    if not req.get("dialog_id"):
-        return get_data_error_result(retmsg="assistant_id is required.")
+    req["dialog_id"] = chat_id
     dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
     if not dia:
-        return get_data_error_result(retmsg="You do not own the assistant")
+        return get_error_data_result(retmsg="You do not own the assistant")
     conv = {
         "id": get_uuid(),
         "dialog_id": req["dialog_id"],
         "name": req.get("name", "New session"),
         "message": [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
     }
     if not conv.get("name"):
-        return get_data_error_result(retmsg="name can not be empty.")
+        return get_error_data_result(retmsg="Name can not be empty.")
     ConversationService.save(**conv)
     e, conv = ConversationService.get_by_id(conv["id"])
     if not e:
-        return get_data_error_result(retmsg="Fail to new session!")
+        return get_error_data_result(retmsg="Fail to create a session!")
     conv = conv.to_dict()
     conv['messages'] = conv.pop("message")
-    conv["assistant_id"] = conv.pop("dialog_id")
+    conv["chat_id"] = conv.pop("dialog_id")
     del conv["reference"]
-    return get_json_result(data=conv)
+    return get_result(data=conv)
 
-
-@manager.route('/completion', methods=['POST'])
+@manager.route('/chat/<chat_id>/session/<session_id>', methods=['PUT'])
 @token_required
-def completion(tenant_id):
+def update(tenant_id,chat_id,session_id):
+    req = request.json
+    if "dialog_id" in req and req.get("dialog_id") != chat_id:
+        return get_error_data_result(retmsg="Can't change chat_id")
+    if "chat_id" in req and req.get("chat_id") != chat_id:
+        return get_error_data_result(retmsg="Can't change chat_id")
+    req["dialog_id"] = chat_id
+    conv_id = session_id
+    conv = ConversationService.query(id=conv_id,dialog_id=chat_id)
+    if not conv:
+        return get_error_data_result(retmsg="Session does not exist")
+    if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+        return get_error_data_result(retmsg="You do not own the session")
+    if "message" in req or "messages" in req:
+        return get_error_data_result(retmsg="Message can not be change")
+    if "reference" in req:
+        return get_error_data_result(retmsg="Reference can not be change")
+    if "name" in req and not req.get("name"):
+        return get_error_data_result(retmsg="Name can not be empty.")
+    if not ConversationService.update_by_id(conv_id, req):
+        return get_error_data_result(retmsg="Session updates error")
+    return get_result()
+
+
+@manager.route('/chat/<chat_id>/session/<session_id>/completion', methods=['POST'])
+@token_required
+def completion(tenant_id,chat_id,session_id):
     req = request.json
     # req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
     #    {"role": "user", "content": "上海有吗？"}
     # ]}
-    if "session_id" not in req:
-        return get_data_error_result(retmsg="session_id is required")
-    conv = ConversationService.query(id=req["session_id"])
+    if not req.get("question"):
+        return get_error_data_result(retmsg="Please input your question.")
+    conv = ConversationService.query(id=session_id,dialog_id=chat_id)
     if not conv:
-        return get_data_error_result(retmsg="Session does not exist")
+        return get_error_data_result(retmsg="Session does not exist")
     conv = conv[0]
-    if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
-        return get_data_error_result(retmsg="You do not own the session")
+    if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+        return get_error_data_result(retmsg="You do not own the session")
     msg = []
     question = {
         "content": req.get("question"),
@@ -108,7 +104,6 @@ def completion(tenant_id):
         msg.append(m)
     message_id = msg[-1].get("id")
     e, dia = DialogService.get_by_id(conv.dialog_id)
-    del req["session_id"]
 
     if not conv.reference:
         conv.reference = []
@@ -130,13 +125,13 @@ def stream():
         try:
             for ans in chat(dia, msg, **req):
                 fillin_conv(ans)
-                yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
+                yield "data:" + json.dumps({"code": 0,  "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),
+            yield "data:" + json.dumps({"code": 500, "message": 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"
+        yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
 
     if req.get("stream", True):
         resp = Response(stream(), mimetype="text/event-stream")
@@ -153,73 +148,31 @@ def stream():
             fillin_conv(ans)
             ConversationService.update_by_id(conv.id, conv.to_dict())
             break
-        return get_json_result(data=answer)
-
+        return get_result(data=answer)
 
-@manager.route('/get', methods=['GET'])
+@manager.route('/chat/<chat_id>/session', methods=['GET'])
 @token_required
-def get(tenant_id):
-    req = request.args
-    if "id" not in req:
-        return get_data_error_result(retmsg="id is required")
-    conv_id = req["id"]
-    conv = ConversationService.query(id=conv_id)
-    if not conv:
-        return get_data_error_result(retmsg="Session does not exist")
-    if not DialogService.query(id=conv[0].dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
-        return get_data_error_result(retmsg="You do not own the session")
-    if "assistant_id" in req:
-        if req["assistant_id"] != conv[0].dialog_id:
-            return get_data_error_result(retmsg="The session doesn't belong to the assistant")
-    conv = conv[0].to_dict()
-    conv['messages'] = conv.pop("message")
-    conv["assistant_id"] = conv.pop("dialog_id")
-    if conv["reference"]:
-        messages = conv["messages"]
-        message_num = 0
-        chunk_num = 0
-        while message_num < len(messages):
-            if message_num != 0 and messages[message_num]["role"] != "user":
-                chunk_list = []
-                if "chunks" in conv["reference"][chunk_num]:
-                    chunks = conv["reference"][chunk_num]["chunks"]
-                    for chunk in chunks:
-                        new_chunk = {
-                            "id": chunk["chunk_id"],
-                            "content": chunk["content_with_weight"],
-                            "document_id": chunk["doc_id"],
-                            "document_name": chunk["docnm_kwd"],
-                            "knowledgebase_id": chunk["kb_id"],
-                            "image_id": chunk["img_id"],
-                            "similarity": chunk["similarity"],
-                            "vector_similarity": chunk["vector_similarity"],
-                            "term_similarity": chunk["term_similarity"],
-                            "positions": chunk["positions"],
-                        }
-                        chunk_list.append(new_chunk)
-                chunk_num += 1
-                messages[message_num]["reference"] = chunk_list
-            message_num += 1
-    del conv["reference"]
-    return get_json_result(data=conv)
-
-
-@manager.route('/list', methods=["GET"])
-@token_required
-def list(tenant_id):
-    assistant_id = request.args["assistant_id"]
-    if not DialogService.query(tenant_id=tenant_id, id=assistant_id, status=StatusEnum.VALID.value):
-        return get_json_result(
-            data=False, retmsg=f"You don't own the assistant.",
-            retcode=RetCode.OPERATING_ERROR)
-    convs = ConversationService.query(
-        dialog_id=assistant_id,
-        order_by=ConversationService.model.create_time,
-        reverse=True)
-    convs = [d.to_dict() for d in convs]
+def list(chat_id,tenant_id):
+    if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
+        return get_error_data_result(retmsg=f"You don't own the assistant {chat_id}.")
+    id = request.args.get("id")
+    name = request.args.get("name")
+    session = ConversationService.query(id=id,name=name,dialog_id=chat_id)
+    if not session:
+        return get_error_data_result(retmsg="The session doesn't exist")
+    page_number = int(request.args.get("page", 1))
+    items_per_page = int(request.args.get("page_size", 1024))
+    orderby = request.args.get("orderby", "create_time")
+    if request.args.get("desc") == "False":
+        desc = False
+    else:
+        desc = True
+    convs = ConversationService.get_list(chat_id,page_number,items_per_page,orderby,desc,id,name)
+    if not convs:
+        return get_result(data=[])
     for conv in convs:
         conv['messages'] = conv.pop("message")
-        conv["assistant_id"] = conv.pop("dialog_id")
+        conv["chat"] = conv.pop("dialog_id")
         if conv["reference"]:
             messages = conv["messages"]
             message_num = 0
@@ -247,20 +200,19 @@ def list(tenant_id):
                     messages[message_num]["reference"] = chunk_list
                 message_num += 1
         del conv["reference"]
-    return get_json_result(data=convs)
+    return get_result(data=convs)
 
-
-@manager.route('/delete', methods=["DELETE"])
+@manager.route('/chat/<chat_id>/session', methods=["DELETE"])
 @token_required
-def delete(tenant_id):
-    id = request.args.get("id")
-    if not id:
-        return get_data_error_result(retmsg="`id` is required in deleting operation")
-    conv = ConversationService.query(id=id)
-    if not conv:
-        return get_data_error_result(retmsg="Session doesn't exist")
-    conv = conv[0]
-    if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
-        return get_data_error_result(retmsg="You don't own the session")
-    ConversationService.delete_by_id(id)
-    return get_json_result(data=True)
+def delete(tenant_id,chat_id):
+    if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+        return get_error_data_result(retmsg="You don't own the chat")
+    ids = request.json.get("ids")
+    if not ids:
+        return get_error_data_result(retmsg="`ids` is required in deleting operation")
+    for id in ids:
+        conv = ConversationService.query(id=id,dialog_id=chat_id)
+        if not conv:
+            return get_error_data_result(retmsg="The chat doesn't own the session")
+        ConversationService.delete_by_id(id)
+    return get_result()

api/apps/sdk/session.py

@@ -19,6 +19,8 @@
 import re
 from copy import deepcopy
 from timeit import default_timer as timer
+
+
 from api.db import LLMType, ParserType,StatusEnum
 from api.db.db_models import Dialog, Conversation,DB
 from api.db.services.common_service import CommonService
@@ -61,6 +63,22 @@ def get_list(cls, tenant_id,
 class ConversationService(CommonService):
     model = Conversation
 
+    @classmethod
+    @DB.connection_context()
+    def get_list(cls,dialog_id,page_number, items_per_page, orderby, desc, id , name):
+        sessions = cls.model.select().where(cls.model.dialog_id ==dialog_id)
+        if id:
+            sessions = sessions.where(cls.model.id == id)
+        if name:
+            sessions = sessions.where(cls.model.name == name)
+        if desc:
+            sessions = sessions.order_by(cls.model.getter_by(orderby).desc())
+        else:
+            sessions = sessions.order_by(cls.model.getter_by(orderby).asc())
+
+        sessions = sessions.paginate(page_number, items_per_page)
+
+        return list(sessions.dicts())
 
 def message_fit_in(msg, max_length=4000):
     def count():

api/db/services/dialog_service.py

@@ -49,6 +49,29 @@
 class DocumentService(CommonService):
     model = Document
 
+    @classmethod
+    @DB.connection_context()
+    def get_list(cls, kb_id, page_number, items_per_page,
+                     orderby, desc, keywords, id):
+        docs =cls.model.select().where(cls.model.kb_id==kb_id)
+        if id:
+            docs = docs.where(
+                cls.model.id== id )
+        if keywords:
+            docs = docs.where(
+                fn.LOWER(cls.model.name).contains(keywords.lower())
+            )
+        count = docs.count()
+        if desc:
+            docs = docs.order_by(cls.model.getter_by(orderby).desc())
+        else:
+            docs = docs.order_by(cls.model.getter_by(orderby).asc())
+
+        docs = docs.paginate(page_number, items_per_page)
+
+        return list(docs.dicts()), count
+
+
     @classmethod
     @DB.connection_context()
     def get_by_kb_id(cls, kb_id, page_number, items_per_page,
@@ -268,7 +291,7 @@ def get_doc_id_by_doc_name(cls, doc_name):
     @classmethod
     @DB.connection_context()
     def get_thumbnails(cls, docids):
-        fields = [cls.model.id, cls.model.kb_id, cls.model.thumbnail]
+        fields = [cls.model.id, cls.model.thumbnail]
         return list(cls.model.select(
             *fields).where(cls.model.id.in_(docids)).dicts())