Merge pull request #16 from ConnectAI-E/feature-lloyd-know

支持跨知识库查询
ConnectAI-E · Nov 28, 2023 · 135247d · 135247d
2 parents e30904b + 0a5a54d
commit 135247d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/server/models.py b/server/models.py
@@ -401,8 +401,9 @@ def get_bot_by_hash(hash):
 def query_by_collection_id(collection_id, q, page, size, delta=None):
     from tasks import embed_query
     embed = embed_query(q)
+    # collection_id支持传数组或者字符串，使用数组的时候，表示跨知识库查询
     filter_ = [{
-        "term": { "collection_id": collection_id },
+        "terms": { "collection_id": collection_id if isinstance(collection_id, list) else [collection_id] },
     }, {
         "term": { "status": 0 },
     }]
@@ -466,11 +467,13 @@ def _query_by_filter_and_embed(q, filter_, embed, page, size, delta=0.5):
         match = match[0].value if len(match) > 0 else 0
         min = match if match < min else min
         max = match if match > max else max
-        explanation.append({
-            'id': i.meta.id,
-            'topk': topk,
-            'match': match,
-        })
+        # 只需要有匹配的，没有任何匹配，就认为完全无关
+        if match > 0:
+            explanation.append({
+                'id': i.meta.id,
+                'topk': topk,
+                'match': match,
+            })
 
     for exp in explanation:
         exp['stand_score'] = (exp['match'] - min) / ((max - min) or 1)

diff --git a/server/routes.py b/server/routes.py
@@ -579,8 +579,11 @@ def api_query_by_collection_id(collection_id):
     size = request.args.get('size', default=20, type=int)
     size = 10000 if size > 10000 else size
     user_id = session.get('user_id', '')
-    collection = get_collection_by_id(user_id, collection_id)
-    assert collection, '找不到知识库或者没有权限'
+    # using array
+    collection_id = collection_id.split(',')
+    for cid in collection_id:
+        collection = get_collection_by_id(user_id, cid)
+        assert collection, '找不到知识库或者没有权限'
 
     documents, total = query_by_collection_id(collection_id, q, page, size)