add conversation API (#35)

api/apps/chunk_app.py

@@ -13,17 +13,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import hashlib
-import re
 
-import numpy as np
 from flask import request
 from flask_login import login_required, current_user
-
-from rag.nlp import search, huqie
+from elasticsearch_dsl import Q
+from rag.nlp import search, huqie, retrievaler
 from rag.utils import ELASTICSEARCH, rmSpace
 from api.db import LLMType
-from api.db.services import duplicate_name
 from api.db.services.kb_service import KnowledgebaseService
 from api.db.services.llm_service import TenantLLMService
 from api.db.services.user_service import UserTenantService
@@ -31,8 +27,9 @@ from api.utils.api_utils import server_error_response, get_data_error_result, va
 from api.db.services.document_service import DocumentService
 from api.settings import RetCode
 from api.utils.api_utils import get_json_result
+import hashlib
+import re
 
-retrival = search.Dealer(ELASTICSEARCH)
 
 @manager.route('/list', methods=['POST'])
 @login_required
@@ -45,12 +42,14 @@ def list():
     question = req.get("keywords", "")
     try:
         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
-        if not tenant_id: return get_data_error_result(retmsg="Tenant not found!")
+        if not tenant_id:
+            return get_data_error_result(retmsg="Tenant not found!")
         query = {
             "doc_ids": [doc_id], "page": page, "size": size, "question": question
         }
-        if "available_int" in req: query["available_int"] = int(req["available_int"])
-        sres = retrival.search(query, search.index_name(tenant_id))
+        if "available_int" in req:
+            query["available_int"] = int(req["available_int"])
+        sres = retrievaler.search(query, search.index_name(tenant_id))
         res = {"total": sres.total, "chunks": []}
         for id in sres.ids:
             d = {
@@ -67,7 +66,7 @@ def list():
     except Exception as e:
         if str(e).find("not_found") > 0:
             return get_json_result(data=False, retmsg=f'Index not found!',
-                            retcode=RetCode.DATA_ERROR)
+                                   retcode=RetCode.DATA_ERROR)
         return server_error_response(e)
 
 
@@ -79,8 +78,11 @@ def get():
         tenants = UserTenantService.query(user_id=current_user.id)
         if not tenants:
             return get_data_error_result(retmsg="Tenant not found!")
-        res = ELASTICSEARCH.get(chunk_id, search.index_name(tenants[0].tenant_id))
-        if not res.get("found"):return server_error_response("Chunk not found")
+        res = ELASTICSEARCH.get(
+            chunk_id, search.index_name(
+                tenants[0].tenant_id))
+        if not res.get("found"):
+            return server_error_response("Chunk not found")
         id = res["_id"]
         res = res["_source"]
         res["chunk_id"] = id
@@ -90,7 +92,8 @@ def get():
                 k.append(n)
             if re.search(r"(_tks|_ltks)", n):
                 res[n] = rmSpace(res[n])
-        for n in k: del res[n]
+        for n in k:
+            del res[n]
 
         return get_json_result(data=res)
     except Exception as e:
@@ -102,7 +105,8 @@ def get():
 
 @manager.route('/set', methods=['POST'])
 @login_required
-@validate_request("doc_id", "chunk_id", "content_ltks", "important_kwd", "docnm_kwd")
+@validate_request("doc_id", "chunk_id", "content_ltks",
+                  "important_kwd")
 def set():
     req = request.json
     d = {"id": req["chunk_id"]}
@@ -110,15 +114,21 @@ def set():
     d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
     d["important_kwd"] = req["important_kwd"]
     d["important_tks"] = huqie.qie(" ".join(req["important_kwd"]))
-    if "available_int" in req: d["available_int"] = req["available_int"]
+    if "available_int" in req:
+        d["available_int"] = req["available_int"]
 
     try:
         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
-        if not tenant_id: return get_data_error_result(retmsg="Tenant not found!")
-        embd_mdl = TenantLLMService.model_instance(tenant_id, LLMType.EMBEDDING.value)
-        v, c = embd_mdl.encode([req["docnm_kwd"], req["content_ltks"]])
+        if not tenant_id:
+            return get_data_error_result(retmsg="Tenant not found!")
+        embd_mdl = TenantLLMService.model_instance(
+            tenant_id, LLMType.EMBEDDING.value)
+        e, doc = DocumentService.get_by_id(req["doc_id"])
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        v, c = embd_mdl.encode([doc.name, req["content_ltks"]])
         v = 0.1 * v[0] + 0.9 * v[1]
-        d["q_%d_vec"%len(v)] = v.tolist()
+        d["q_%d_vec" % len(v)] = v.tolist()
         ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
         return get_json_result(data=True)
     except Exception as e:
@@ -132,19 +142,32 @@ def switch():
     req = request.json
     try:
         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
-        if not tenant_id: return get_data_error_result(retmsg="Tenant not found!")
+        if not tenant_id:
+            return get_data_error_result(retmsg="Tenant not found!")
         if not ELASTICSEARCH.upsert([{"id": i, "available_int": int(req["available_int"])} for i in req["chunk_ids"]],
-                             search.index_name(tenant_id)):
+                                    search.index_name(tenant_id)):
             return get_data_error_result(retmsg="Index updating failure")
         return get_json_result(data=True)
     except Exception as e:
         return server_error_response(e)
 
 
+@manager.route('/rm', methods=['POST'])
+@login_required
+@validate_request("chunk_ids")
+def rm():
+    req = request.json
+    try:
+        if not ELASTICSEARCH.deleteByQuery(Q("ids", values=req["chunk_ids"]), search.index_name(current_user.id)):
+            return get_data_error_result(retmsg="Index updating failure")
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)
+
 
 @manager.route('/create', methods=['POST'])
 @login_required
-@validate_request("doc_id", "content_ltks", "important_kwd")
+@validate_request("doc_id", "content_ltks")
 def create():
     req = request.json
     md5 = hashlib.md5()
@@ -152,24 +175,27 @@ def create():
     chunck_id = md5.hexdigest()
     d = {"id": chunck_id, "content_ltks": huqie.qie(req["content_ltks"])}
     d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
-    d["important_kwd"] = req["important_kwd"]
-    d["important_tks"] = huqie.qie(" ".join(req["important_kwd"]))
+    d["important_kwd"] = req.get("important_kwd", [])
+    d["important_tks"] = huqie.qie(" ".join(req.get("important_kwd", [])))
 
     try:
         e, doc = DocumentService.get_by_id(req["doc_id"])
-        if not e: return get_data_error_result(retmsg="Document not found!")
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
         d["kb_id"] = [doc.kb_id]
         d["docnm_kwd"] = doc.name
         d["doc_id"] = doc.id
 
         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
-        if not tenant_id: return get_data_error_result(retmsg="Tenant not found!")
+        if not tenant_id:
+            return get_data_error_result(retmsg="Tenant not found!")
 
-        embd_mdl = TenantLLMService.model_instance(tenant_id, LLMType.EMBEDDING.value)
+        embd_mdl = TenantLLMService.model_instance(
+            tenant_id, LLMType.EMBEDDING.value)
         v, c = embd_mdl.encode([doc.name, req["content_ltks"]])
         DocumentService.increment_chunk_num(req["doc_id"], doc.kb_id, c, 1, 0)
         v = 0.1 * v[0] + 0.9 * v[1]
-        d["q_%d_vec"%len(v)] = v.tolist()
+        d["q_%d_vec" % len(v)] = v.tolist()
         ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
         return get_json_result(data={"chunk_id": chunck_id})
     except Exception as e:
@@ -194,44 +220,15 @@ def retrieval_test():
         if not e:
             return get_data_error_result(retmsg="Knowledgebase not found!")
 
-        embd_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.EMBEDDING.value)
-        sres = retrival.search({"kb_ids": [kb_id], "doc_ids": doc_ids, "size": top,
-                                "question": question, "vector": True,
-                                "similarity": similarity_threshold},
-                               search.index_name(kb.tenant_id),
-                               embd_mdl)
-
-        sim, tsim, vsim = retrival.rerank(sres, question, 1-vector_similarity_weight, vector_similarity_weight)
-        idx = np.argsort(sim*-1)
-        ranks = {"total": 0, "chunks": [], "doc_aggs": {}}
-        start_idx = (page-1)*size
-        for i in idx:
-            ranks["total"] += 1
-            if sim[i] < similarity_threshold: break
-            start_idx -= 1
-            if start_idx >= 0:continue
-            if len(ranks["chunks"]) == size:continue
-            id = sres.ids[i]
-            dnm = sres.field[id]["docnm_kwd"]
-            d = {
-                "chunk_id": id,
-                "content_ltks": sres.field[id]["content_ltks"],
-                "doc_id": sres.field[id]["doc_id"],
-                "docnm_kwd": dnm,
-                "kb_id": sres.field[id]["kb_id"],
-                "important_kwd": sres.field[id].get("important_kwd", []),
-                "img_id": sres.field[id].get("img_id", ""),
-                "similarity": sim[i],
-                "vector_similarity": vsim[i],
-                "term_similarity": tsim[i]
-            }
-            ranks["chunks"].append(d)
-            if dnm not in ranks["doc_aggs"]:ranks["doc_aggs"][dnm] = 0
-            ranks["doc_aggs"][dnm] += 1
+        embd_mdl = TenantLLMService.model_instance(
+            kb.tenant_id, LLMType.EMBEDDING.value)
+        ranks = retrievaler.retrieval(question, embd_mdl, kb.tenant_id, [kb_id], page, size, similarity_threshold,
+                          vector_similarity_weight, top, doc_ids)
 
         return get_json_result(data=ranks)
     except Exception as e:
         if str(e).find("not_found") > 0:
             return get_json_result(data=False, retmsg=f'Index not found!',
-                            retcode=RetCode.DATA_ERROR)
-        return server_error_response(e)
+                                   retcode=RetCode.DATA_ERROR)
+        return server_error_response(e)
+

api/apps/conversation_app.py

@@ -0,0 +1,207 @@
+#
+#  Copyright 2019 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.
+#
+import re
+
+import tiktoken
+from flask import request
+from flask_login import login_required, current_user
+from api.db.services.dialog_service import DialogService, ConversationService
+from api.db import StatusEnum, LLMType
+from api.db.services.kb_service import KnowledgebaseService
+from api.db.services.llm_service import LLMService, TenantLLMService
+from api.db.services.user_service import TenantService
+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
+from rag.llm import ChatModel
+from rag.nlp import retrievaler
+from rag.nlp.query import EsQueryer
+from rag.utils import num_tokens_from_string, encoder
+
+
+@manager.route('/set', methods=['POST'])
+@login_required
+@validate_request("dialog_id")
+def set():
+    req = request.json
+    conv_id = req.get("conversation_id")
+    if conv_id:
+        del req["conversation_id"]
+        try:
+            if not ConversationService.update_by_id(conv_id, req):
+                return get_data_error_result(retmsg="Conversation not found!")
+            e, conv = ConversationService.get_by_id(conv_id)
+            if not e:
+                return get_data_error_result(
+                    retmsg="Fail to update a conversation!")
+            conv = conv.to_dict()
+            return get_json_result(data=conv)
+        except Exception as e:
+            return server_error_response(e)
+
+    try:
+        e, dia = DialogService.get_by_id(req["dialog_id"])
+        if not e:
+            return get_data_error_result(retmsg="Dialog not found")
+        conv = {
+            "id": get_uuid(),
+            "dialog_id": req["dialog_id"],
+            "name": "New conversation",
+            "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
+        }
+        ConversationService.save(**conv)
+        e, conv = ConversationService.get_by_id(conv["id"])
+        if not e:
+            return get_data_error_result(retmsg="Fail to new a conversation!")
+        conv = conv.to_dict()
+        return get_json_result(data=conv)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/get', methods=['GET'])
+@login_required
+def get():
+    conv_id = request.args["conversation_id"]
+    try:
+        e, conv = ConversationService.get_by_id(conv_id)
+        if not e:
+            return get_data_error_result(retmsg="Conversation not found!")
+        conv = conv.to_dict()
+        return get_json_result(data=conv)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/rm', methods=['POST'])
+@login_required
+def rm():
+    conv_ids = request.json["conversation_ids"]
+    try:
+        for cid in conv_ids:
+            ConversationService.delete_by_id(cid)
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)
+
+@manager.route('/list', methods=['GET'])
+@login_required
+def list():
+    dialog_id = request.args["dialog_id"]
+    try:
+        convs = ConversationService.query(dialog_id=dialog_id)
+        convs = [d.to_dict() for d in convs]
+        return get_json_result(data=convs)
+    except Exception as e:
+        return server_error_response(e)
+
+
+def message_fit_in(msg, max_length=4000):
+    def count():
+        nonlocal msg
+        tks_cnts = []
+        for m in msg:tks_cnts.append({"role": m["role"], "count": num_tokens_from_string(m["content"])})
+        total = 0
+        for m in tks_cnts: total += m["count"]
+        return total
+
+    c = count()
+    if c < max_length: return c, msg
+    msg = [m for m in msg if m.role in ["system", "user"]]
+    c = count()
+    if c < max_length:return c, msg
+    msg_ = [m for m in msg[:-1] if m.role == "system"]
+    msg_.append(msg[-1])
+    msg = msg_
+    c = count()
+    if c < max_length:return c, msg
+    ll = num_tokens_from_string(msg_[0].content)
+    l = num_tokens_from_string(msg_[-1].content)
+    if ll/(ll + l) > 0.8:
+        m = msg_[0].content
+        m = encoder.decode(encoder.encode(m)[:max_length-l])
+        msg[0].content = m
+        return max_length, msg
+
+    m = msg_[1].content
+    m = encoder.decode(encoder.encode(m)[:max_length-l])
+    msg[1].content = m
+    return max_length, msg
+
+
+def chat(dialog, messages, **kwargs):
+    assert messages[-1]["role"] == "user", "The last content of this conversation is not from user."
+    llm = LLMService.query(llm_name=dialog.llm_id)
+    if not llm:
+        raise LookupError("LLM(%s) not found"%dialog.llm_id)
+    llm = llm[0]
+    prompt_config = dialog.prompt_config
+    for p in prompt_config["parameters"]:
+        if p["key"] == "knowledge":continue
+        if p["key"] not in kwargs and not p["optional"]:raise KeyError("Miss parameter: " + p["key"])
+        if p["key"] not in kwargs:
+            prompt_config["system"] = prompt_config["system"].replace("{%s}"%p["key"], " ")
+
+    model_config = TenantLLMService.get_api_key(dialog.tenant_id, LLMType.CHAT.value, dialog.llm_id)
+    if not model_config: raise LookupError("LLM(%s) API key not found"%dialog.llm_id)
+
+    question = messages[-1]["content"]
+    embd_mdl = TenantLLMService.model_instance(
+        dialog.tenant_id, LLMType.EMBEDDING.value)
+    kbinfos = retrievaler.retrieval(question, embd_mdl, dialog.tenant_id, dialog.kb_ids, 1, dialog.top_n, dialog.similarity_threshold,
+                        dialog.vector_similarity_weight, top=1024, aggs=False)
+    knowledges = [ck["content_ltks"] for ck in kbinfos["chunks"]]
+
+    if not knowledges and prompt_config["empty_response"]:
+        return {"answer": prompt_config["empty_response"], "retrieval": kbinfos}
+
+    kwargs["knowledge"] = "\n".join(knowledges)
+    gen_conf = dialog.llm_setting[dialog.llm_setting_type]
+    msg = [{"role": m["role"], "content": m["content"]} for m in messages if m["role"] != "system"]
+    used_token_count = message_fit_in(msg, int(llm.max_tokens * 0.97))
+    if "max_tokens" in gen_conf:
+        gen_conf["max_tokens"] = min(gen_conf["max_tokens"], llm.max_tokens - used_token_count)
+    mdl = ChatModel[model_config.llm_factory](model_config["api_key"], dialog.llm_id)
+    answer = mdl.chat(prompt_config["system"].format(**kwargs), msg, gen_conf)
+
+    answer = retrievaler.insert_citations(answer,
+                                 [ck["content_ltks"] for ck in kbinfos["chunks"]],
+                                 [ck["vector"] for ck in kbinfos["chunks"]],
+                                 embd_mdl,
+                                 tkweight=1-dialog.vector_similarity_weight,
+                                 vtweight=dialog.vector_similarity_weight)
+    return {"answer": answer, "retrieval": kbinfos}
+
+
+@manager.route('/completion', methods=['POST'])
+@login_required
+@validate_request("dialog_id", "messages")
+def completion():
+    req = request.json
+    msg = []
+    for m in req["messages"]:
+        if m["role"] == "system":continue
+        if m["role"] == "assistant" and not msg:continue
+        msg.append({"role": m["role"], "content": m["content"]})
+    try:
+        e, dia = DialogService.get_by_id(req["dialog_id"])
+        if not e:
+            return get_data_error_result(retmsg="Dialog not found!")
+        del req["dialog_id"]
+        del req["messages"]
+        return get_json_result(data=chat(dia, msg, **req))
+    except Exception as e:
+        return server_error_response(e)

api/apps/dialog_app.py

@@ -13,28 +13,16 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import hashlib
-import re
 
-import numpy as np
 from flask import request
 from flask_login import login_required, current_user
-
 from api.db.services.dialog_service import DialogService
-from rag.nlp import search, huqie
-from rag.utils import ELASTICSEARCH, rmSpace
-from api.db import LLMType, StatusEnum
-from api.db.services import duplicate_name
+from api.db import StatusEnum
 from api.db.services.kb_service import KnowledgebaseService
-from api.db.services.llm_service import TenantLLMService
-from api.db.services.user_service import UserTenantService, TenantService
+from api.db.services.user_service import TenantService
 from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 from api.utils import get_uuid
-from api.db.services.document_service import DocumentService
-from api.settings import RetCode, stat_logger
 from api.utils.api_utils import get_json_result
-from rag.utils.minio_conn import MINIO
-from api.utils.file_utils import filename_type
 
 
 @manager.route('/set', methods=['POST'])
@@ -128,6 +116,7 @@ def set():
     except Exception as e:
         return server_error_response(e)
 
+
 @manager.route('/get', methods=['GET'])
 @login_required
 def get():
@@ -159,5 +148,18 @@ def list():
         for d in diags:
             d["kb_ids"], d["kb_names"] = get_kb_names(d["kb_ids"])
         return get_json_result(data=diags)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/rm', methods=['POST'])
+@login_required
+@validate_request("dialog_id")
+def rm():
+    req = request.json
+    try:
+        if not DialogService.update_by_id(req["dialog_id"], {"status": StatusEnum.INVALID.value}):
+            return get_data_error_result(retmsg="Dialog not found!")
+        return get_json_result(data=True)
     except Exception as e:
         return server_error_response(e)

api/apps/document_app.py

@@ -271,7 +271,7 @@ def change_parser():
 
 
 @manager.route('/image/<image_id>', methods=['GET'])
-@login_required
+#@login_required
 def get_image(image_id):
     try:
         bkt, nm = image_id.split("-")

api/apps/kb_app.py

@@ -108,7 +108,7 @@ def rm():
         if not KnowledgebaseService.query(created_by=current_user.id, id=req["kb_id"]):
             return get_json_result(data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.', retcode=RetCode.OPERATING_ERROR)
 
-        if not KnowledgebaseService.update_by_id(req["kb_id"], {"status": StatusEnum.IN_VALID.value}): return get_data_error_result(retmsg="Database error (Knowledgebase removal)!")
+        if not KnowledgebaseService.update_by_id(req["kb_id"], {"status": StatusEnum.INVALID.value}): return get_data_error_result(retmsg="Database error (Knowledgebase removal)!")
         return get_json_result(data=True)
     except Exception as e:
         return server_error_response(e)