Refactor Chunk API

api/apps/sdk/doc.py

@@ -494,25 +494,24 @@ def set(tenant_id,dataset_id,document_id,chunk_id):
 
 
 
-@manager.route('/retrieval', methods=['GET'])
+@manager.route('/retrieval', methods=['POST'])
 @token_required
 def retrieval_test(tenant_id):
-    req = request.args
-    req_json = request.json
-    if not req_json.get("datasets"):
+    req = request.json
+    if not req.get("datasets"):
         return get_error_data_result("`datasets` is required.")
-    for id in req_json.get("datasets"):
+    kb_id = req["datasets"]
+    if isinstance(kb_id, str): kb_id = [kb_id]
+    for id in kb_id:
         if not KnowledgebaseService.query(id=id,tenant_id=tenant_id):
             return get_error_data_result(f"You don't own the dataset {id}.")
-    if "question" not in req_json:
+    if "question" not in req:
         return get_error_data_result("`question` is required.")
     page = int(req.get("offset", 1))
     size = int(req.get("limit", 30))
-    question = req_json["question"]
-    kb_id = req_json["datasets"]
-    if isinstance(kb_id, str): kb_id = [kb_id]
-    doc_ids = req_json.get("documents", [])
-    similarity_threshold = float(req.get("similarity_threshold", 0.0))
+    question = req["question"]
+    doc_ids = req.get("documents", [])
+    similarity_threshold = float(req.get("similarity_threshold", 0.2))
     vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
     top = int(req.get("top_k", 1024))
     if req.get("highlight")=="False" or  req.get("highlight")=="false":

api/apps/sdk/session.py

@@ -24,10 +24,9 @@
 from api.utils.api_utils import get_error_data_result
 from api.utils.api_utils import get_result, token_required
 
-
 @manager.route('/chat/<chat_id>/session', methods=['POST'])
 @token_required
-def create(tenant_id, chat_id):
+def create(tenant_id,chat_id):
     req = request.json
     req["dialog_id"] = chat_id
     dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
@@ -51,14 +50,13 @@ def create(tenant_id, chat_id):
     del conv["reference"]
     return get_result(data=conv)
 
-
 @manager.route('/chat/<chat_id>/session/<session_id>', methods=['PUT'])
 @token_required
-def update(tenant_id, chat_id, session_id):
+def update(tenant_id,chat_id,session_id):
     req = request.json
     req["dialog_id"] = chat_id
     conv_id = session_id
-    conv = ConversationService.query(id=conv_id, dialog_id=chat_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):
@@ -74,16 +72,30 @@ def update(tenant_id, chat_id, session_id):
     return get_result()
 
 
-@manager.route('/chat/<chat_id>/session/<session_id>/completion', methods=['POST'])
+@manager.route('/chat/<chat_id>/completion', methods=['POST'])
 @token_required
-def completion(tenant_id, chat_id, session_id):
+def completion(tenant_id,chat_id):
     req = request.json
     # req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
     #    {"role": "user", "content": "上海有吗？"}
     # ]}
+    if not req.get("session_id"):
+        conv = {
+            "id": get_uuid(),
+            "dialog_id": chat_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_error_data_result(retmsg="Name can not be empty.")
+        ConversationService.save(**conv)
+        e, conv = ConversationService.get_by_id(conv["id"])
+        session_id=conv.id
+    else:
+        session_id = req.get("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)
+    conv = ConversationService.query(id=session_id,dialog_id=chat_id)
     if not conv:
         return get_error_data_result(retmsg="Session does not exist")
     conv = conv[0]
@@ -123,13 +135,14 @@ def stream():
         try:
             for ans in chat(dia, msg, **req):
                 fillin_conv(ans)
-                yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
+                yield "data:" + json.dumps({"code": 0,  "data": {**ans, "session_id":session_id}}, ensure_ascii=False) + "\n\n"
+
             ConversationService.update_by_id(conv.id, conv.to_dict())
         except Exception as e:
             yield "data:" + json.dumps({"code": 500, "message": str(e),
-                                        "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
+                                        "data": {"answer": "**ERROR**: " + str(e), "session_id":session_id,"reference": []}},
                                        ensure_ascii=False) + "\n\n"
-        yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "data": True,"session_id":session_id}, ensure_ascii=False) + "\n\n"
 
     if req.get("stream", True):
         resp = Response(stream(), mimetype="text/event-stream")
@@ -148,15 +161,14 @@ def stream():
             break
         return get_result(data=answer)
 
-
 @manager.route('/chat/<chat_id>/session', methods=['GET'])
 @token_required
-def list(chat_id, tenant_id):
+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)
+    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))
@@ -166,7 +178,7 @@ def list(chat_id, tenant_id):
         desc = False
     else:
         desc = True
-    convs = ConversationService.get_list(chat_id, page_number, items_per_page, orderby, desc, id, name)
+    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:
@@ -201,17 +213,16 @@ def list(chat_id, tenant_id):
         del conv["reference"]
     return get_result(data=convs)
 
-
 @manager.route('/chat/<chat_id>/session', methods=["DELETE"])
 @token_required
-def delete(tenant_id, chat_id):
+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)
+        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)

api/db/services/document_service.py

@@ -61,14 +61,13 @@ def get_list(cls, kb_id, page_number, items_per_page,
             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)
-
+        count = docs.count()
         return list(docs.dicts()), count
 
 

sdk/python/ragflow/modules/chunk.py

@@ -17,32 +17,11 @@ def __init__(self, rag, res_dict):
                 res_dict.pop(k)
         super().__init__(rag, res_dict)
 
-    def delete(self) -> bool:
-        """
-        Delete the chunk in the document.
-        """
-        res = self.post('/doc/chunk/rm',
-                        {"document_id": self.document_id, 'chunk_ids': [self.id]})
-        res = res.json()
-        if res.get("retmsg") == "success":
-            return True
-        raise Exception(res["retmsg"])
 
-    def save(self) -> bool:
-        """
-        Save the document details to the server.
-        """
-        res = self.post('/doc/chunk/set',
-                        {"chunk_id": self.id,
-                         "knowledgebase_id": self.knowledgebase_id,
-                         "name": self.document_name,
-                         "content": self.content,
-                         "important_keywords": self.important_keywords,
-                         "document_id": self.document_id,
-                         "available": self.available,
-                         })
+    def update(self,update_message:dict):
+        res = self.put(f"/dataset/{self.knowledgebase_id}/document/{self.document_id}/chunk/{self.id}",update_message)
         res = res.json()
-        if res.get("retmsg") == "success":
-            return True
-        raise Exception(res["retmsg"])
+        if res.get("code") != 0 :
+            raise Exception(res["message"])
+
 

sdk/python/ragflow/modules/dataset.py

@@ -65,3 +65,14 @@ def delete_documents(self,ids: List[str] = None):
         if res.get("code") != 0:
             raise Exception(res["message"])
 
+    def async_parse_documents(self,document_ids):
+        res = self.post(f"/dataset/{self.id}/chunk",{"document_ids":document_ids})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))
+
+    def async_cancel_parse_documents(self,document_ids):
+        res = self.rm(f"/dataset/{self.id}/chunk",{"document_ids":document_ids})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))