change licence (#28)

* add front end code

* change licence

rag/llm/__init__.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2019 The FATE Authors. All Rights Reserved.
+#  Copyright 2019 The RAG Flow 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.

rag/llm/chat_model.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2019 The FATE Authors. All Rights Reserved.
+#  Copyright 2019 The RAG Flow 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.

rag/llm/cv_model.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2019 The FATE Authors. All Rights Reserved.
+#  Copyright 2019 The RAG Flow 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.

rag/llm/embedding_model.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2019 The FATE Authors. All Rights Reserved.
+#  Copyright 2019 The RAG Flow 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.

rag/nlp/search.py

@@ -1,8 +1,11 @@
 # -*- coding: utf-8 -*-
+import json
 import re
 from elasticsearch_dsl import Q, Search, A
 from typing import List, Optional, Tuple, Dict, Union
 from dataclasses import dataclass
+
+from rag.settings import es_logger
 from rag.utils import rmSpace
 from rag.nlp import huqie, query
 import numpy as np
@@ -34,30 +37,30 @@ class Dealer:
         group_docs: List[List] = None
 
     def _vector(self, txt, sim=0.8, topk=10):
+        qv = self.emb_mdl.encode_queries(txt)
         return {
-            "field": "q_vec",
+            "field": "q_%d_vec"%len(qv),
             "k": topk,
             "similarity": sim,
             "num_candidates": 1000,
-            "query_vector": self.emb_mdl.encode_queries(txt)
+            "query_vector": qv
         }
 
     def search(self, req, idxnm, tks_num=3):
-        keywords = []
         qst = req.get("question", "")
-
         bqry, keywords = self.qryr.question(qst)
         if req.get("kb_ids"):
             bqry.filter.append(Q("terms", kb_id=req["kb_ids"]))
-        bqry.filter.append(Q("exists", field="q_tks"))
+        if req.get("doc_ids"):
+            bqry.filter.append(Q("terms", doc_id=req["doc_ids"]))
         bqry.boost = 0.05
-        print(bqry)
 
         s = Search()
         pg = int(req.get("page", 1)) - 1
         ps = int(req.get("size", 1000))
-        src = req.get("field", ["docnm_kwd", "content_ltks", "kb_id",
-                                "image_id", "doc_id", "q_vec"])
+        src = req.get("fields", ["docnm_kwd", "content_ltks", "kb_id","img_id",
+                                "image_id", "doc_id", "q_512_vec", "q_768_vec",
+                                "q_1024_vec", "q_1536_vec"])
 
         s = s.query(bqry)[pg * ps:(pg + 1) * ps]
         s = s.highlight("content_ltks")
@@ -66,22 +69,24 @@ class Dealer:
             s = s.sort(
                 {"create_time": {"order": "desc", "unmapped_type": "date"}})
 
-        s = s.highlight_options(
-            fragment_size=120,
-            number_of_fragments=5,
-            boundary_scanner_locale="zh-CN",
-            boundary_scanner="SENTENCE",
-            boundary_chars=",./;:\\!()，。？：！……（）——、"
-        )
+        if qst:
+            s = s.highlight_options(
+                fragment_size=120,
+                number_of_fragments=5,
+                boundary_scanner_locale="zh-CN",
+                boundary_scanner="SENTENCE",
+                boundary_chars=",./;:\\!()，。？：！……（）——、"
+            )
         s = s.to_dict()
         q_vec = []
         if req.get("vector"):
             s["knn"] = self._vector(qst, req.get("similarity", 0.4), ps)
             s["knn"]["filter"] = bqry.to_dict()
-            del s["highlight"]
+            if "highlight" in s: del s["highlight"]
             q_vec = s["knn"]["query_vector"]
+        es_logger.info("【Q】: {}".format(json.dumps(s)))
         res = self.es.search(s, idxnm=idxnm, timeout="600s", src=src)
-        print("TOTAL: ", self.es.getTotal(res))
+        es_logger.info("TOTAL: {}".format(self.es.getTotal(res)))
         if self.es.getTotal(res) == 0 and "knn" in s:
             bqry, _ = self.qryr.question(qst, min_match="10%")
             if req.get("kb_ids"):
@@ -109,8 +114,7 @@ class Dealer:
             query_vector=q_vec,
             aggregation=aggs,
             highlight=self.getHighlight(res),
-            field=self.getFields(res, ["docnm_kwd", "content_ltks",
-                                       "kb_id", "image_id", "doc_id", "q_vec"]),
+            field=self.getFields(res, src),
             keywords=list(kwds)
         )
 
@@ -237,14 +241,4 @@ class Dealer:
         return sim
 
 
-if __name__ == "__main__":
-    from util import es_conn
-    SE = Dealer(es_conn.HuEs("infiniflow"))
-    qs = [
-        "胡凯",
-        ""
-    ]
-    for q in qs:
-        print(">>>>>>>>>>>>>>>>>>>>", q)
-        print(SE.search(
-            {"question": q, "kb_ids": "64f072a75f3b97c865718c4a"}, "infiniflow_*"))
+

rag/nlp/term_weight.py

@@ -62,7 +62,7 @@ class Dealer:
                 return set(res.keys())
             return res
 
-        fnm = os.path.join(get_project_base_directory(), "res")
+        fnm = os.path.join(get_project_base_directory(), "rag/res")
         self.ne, self.df = {}, {}
         try:
             self.ne = json.load(open(os.path.join(fnm, "ner.json"), "r"))

rag/settings.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2019 The FATE Authors. All Rights Reserved.
+#  Copyright 2019 The RAG Flow 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.

rag/svr/parse_user_docs.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2019 The FATE Authors. All Rights Reserved.
+#  Copyright 2019 The RAG Flow 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.
@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import datetime
 import json
 import logging
 import os
@@ -108,17 +109,17 @@ def build(row, cvmdl):
                      (int(DOC_MAXIMUM_SIZE / 1024 / 1024)))
         return []
 
-    res = ELASTICSEARCH.search(Q("term", doc_id=row["id"]))
-    if ELASTICSEARCH.getTotal(res) > 0:
-        ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=row["id"]),
-                                          scripts="""
-                               if(!ctx._source.kb_id.contains('%s'))
-                                 ctx._source.kb_id.add('%s');
-                               """ % (str(row["kb_id"]), str(row["kb_id"])),
-            idxnm=search.index_name(row["tenant_id"])
-        )
-        set_progress(row["id"], 1, "Done")
-        return []
+    # res = ELASTICSEARCH.search(Q("term", doc_id=row["id"]))
+    # if ELASTICSEARCH.getTotal(res) > 0:
+    #     ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=row["id"]),
+    #                                       scripts="""
+    #                            if(!ctx._source.kb_id.contains('%s'))
+    #                              ctx._source.kb_id.add('%s');
+    #                            """ % (str(row["kb_id"]), str(row["kb_id"])),
+    #         idxnm=search.index_name(row["tenant_id"])
+    #     )
+    #     set_progress(row["id"], 1, "Done")
+    #     return []
 
     random.seed(time.time())
     set_progress(row["id"], random.randint(0, 20) /
@@ -155,8 +156,7 @@ def build(row, cvmdl):
         "doc_id": row["id"],
         "kb_id": [str(row["kb_id"])],
         "docnm_kwd": os.path.split(row["location"])[-1],
-        "title_tks": huqie.qie(row["name"]),
-        "updated_at": str(row["update_time"]).replace("T", " ")[:19]
+        "title_tks": huqie.qie(row["name"])
     }
     doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
     output_buffer = BytesIO()
@@ -179,6 +179,7 @@ def build(row, cvmdl):
 
         MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue())
         d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"])
+        d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
         docs.append(d)
 
     for arr, img in obj.table_chunks:
@@ -193,6 +194,7 @@ def build(row, cvmdl):
             img.save(output_buffer, format='JPEG')
             MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue())
             d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"])
+            d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
             docs.append(d)
     set_progress(row["id"], random.randint(60, 70) /
                  100., "Continue embedding the content.")
@@ -218,23 +220,11 @@ def embedding(docs, mdl):
     vects = 0.1 * tts + 0.9 * cnts
     assert len(vects) == len(docs)
     for i, d in enumerate(docs):
-        d["q_vec"] = vects[i].tolist()
+        v = vects[i].tolist()
+        d["q_%d_vec"%len(v)] = v
     return tk_count
 
 
-def model_instance(tenant_id, llm_type):
-    model_config = TenantLLMService.get_api_key(tenant_id, model_type=LLMType.EMBEDDING)
-    if not model_config:
-        model_config = {"llm_factory": "local", "api_key": "", "llm_name": ""}
-    else: model_config = model_config[0].to_dict()
-    if llm_type == LLMType.EMBEDDING:
-        if model_config["llm_factory"] not in EmbeddingModel: return
-        return EmbeddingModel[model_config["llm_factory"]](model_config["api_key"], model_config["llm_name"])
-    if llm_type == LLMType.IMAGE2TEXT:
-        if model_config["llm_factory"] not in CvModel: return
-        return CvModel[model_config.llm_factory](model_config["api_key"], model_config["llm_name"])
-
-
 def main(comm, mod):
     global model
     from rag.llm import HuEmbedding
@@ -247,12 +237,12 @@ def main(comm, mod):
 
     tmf = open(tm_fnm, "a+")
     for _, r in rows.iterrows():
-        embd_mdl = model_instance(r["tenant_id"], LLMType.EMBEDDING)
+        embd_mdl = TenantLLMService.model_instance(r["tenant_id"], LLMType.EMBEDDING)
         if not embd_mdl:
             set_progress(r["id"], -1, "Can't find embedding model!")
             cron_logger.error("Tenant({}) can't find embedding model!".format(r["tenant_id"]))
             continue
-        cv_mdl = model_instance(r["tenant_id"], LLMType.IMAGE2TEXT)
+        cv_mdl = TenantLLMService.model_instance(r["tenant_id"], LLMType.IMAGE2TEXT)
         st_tm = timer()
         cks = build(r, cv_mdl)
         if not cks:

rag/utils/es_conn.py

@@ -241,6 +241,26 @@ class HuEs:
         es_logger.error("ES search timeout for 3 times!")
         raise Exception("ES search timeout.")
 
+    def get(self, doc_id, idxnm=None):
+        for i in range(3):
+            try:
+                res = self.es.get(index=(self.idxnm if not idxnm else idxnm),
+                                     id=doc_id)
+                if str(res.get("timed_out", "")).lower() == "true":
+                    raise Exception("Es Timeout.")
+                return res
+            except Exception as e:
+                es_logger.error(
+                    "ES get exception: " +
+                    str(e) +
+                    "【Q】：" +
+                    doc_id)
+                if str(e).find("Timeout") > 0:
+                    continue
+                raise e
+        es_logger.error("ES search timeout for 3 times!")
+        raise Exception("ES search timeout.")
+
     def updateByQuery(self, q, d):
         ubq = UpdateByQuery(index=self.idxnm).using(self.es).query(q)
         scripts = ""