HowardChan commited on
Commit
ebde808
·
1 Parent(s): 342ad40

Feat:Compatible with Dify's External Knowledge API (#2848)

Browse files

### What problem does this PR solve?

_Briefly describe what this PR aims to solve. Include background context
that will help reviewers understand the purpose of the PR._
Fixes #2731
### Type of change

- [x] New Feature (non-breaking change which adds functionality)

api/apps/sdk/dify_retrieval.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from flask import request, jsonify
2
+
3
+ from db import LLMType, ParserType
4
+ from db.services.knowledgebase_service import KnowledgebaseService
5
+ from db.services.llm_service import LLMBundle
6
+ from settings import retrievaler, kg_retrievaler, RetCode
7
+ from utils.api_utils import validate_request, build_error_result, apikey_required
8
+
9
+
10
+ @manager.route('/dify/retrieval', methods=['POST'])
11
+ @apikey_required
12
+ @validate_request("knowledge_id", "query")
13
+ def retrieval(tenant_id):
14
+ req = request.json
15
+ question = req["query"]
16
+ kb_id = req["knowledge_id"]
17
+ retrieval_setting = req.get("retrieval_setting", {})
18
+ similarity_threshold = float(retrieval_setting.get("score_threshold", 0.0))
19
+ top = int(retrieval_setting.get("top_k", 1024))
20
+
21
+ try:
22
+
23
+ e, kb = KnowledgebaseService.get_by_id(kb_id)
24
+ if not e:
25
+ return build_error_result(error_msg="Knowledgebase not found!", retcode=RetCode.NOT_FOUND)
26
+
27
+ if kb.tenant_id != tenant_id:
28
+ return build_error_result(error_msg="Knowledgebase not found!", retcode=RetCode.NOT_FOUND)
29
+
30
+ embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
31
+
32
+ retr = retrievaler if kb.parser_id != ParserType.KG else kg_retrievaler
33
+ ranks = retr.retrieval(
34
+ question,
35
+ embd_mdl,
36
+ kb.tenant_id,
37
+ [kb_id],
38
+ page=1,
39
+ page_size=top,
40
+ similarity_threshold=similarity_threshold,
41
+ vector_similarity_weight=0.3,
42
+ top=top
43
+ )
44
+ records = []
45
+ for c in ranks["chunks"]:
46
+ if "vector" in c:
47
+ del c["vector"]
48
+ records.append({
49
+ "content": c["content_ltks"],
50
+ "score": c["similarity"],
51
+ "title": c["docnm_kwd"],
52
+ "metadata": ""
53
+ })
54
+
55
+ return jsonify({"records": records})
56
+ except Exception as e:
57
+ if str(e).find("not_found") > 0:
58
+ return build_error_result(
59
+ error_msg=f'No chunk found! Check the chunk status please!',
60
+ retcode=RetCode.NOT_FOUND
61
+ )
62
+ return build_error_result(error_msg=str(e), retcode=RetCode.SERVER_ERROR)
api/settings.py CHANGED
@@ -250,3 +250,5 @@ class RetCode(IntEnum, CustomEnum):
250
  AUTHENTICATION_ERROR = 109
251
  UNAUTHORIZED = 401
252
  SERVER_ERROR = 500
 
 
 
250
  AUTHENTICATION_ERROR = 109
251
  UNAUTHORIZED = 401
252
  SERVER_ERROR = 500
253
+ FORBIDDEN = 403
254
+ NOT_FOUND = 404
api/utils/api_utils.py CHANGED
@@ -200,6 +200,27 @@ def get_json_result(retcode=RetCode.SUCCESS, retmsg='success', data=None):
200
  response = {"retcode": retcode, "retmsg": retmsg, "data": data}
201
  return jsonify(response)
202
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
203
 
204
  def construct_response(retcode=RetCode.SUCCESS,
205
  retmsg='success', data=None, auth=None):
 
200
  response = {"retcode": retcode, "retmsg": retmsg, "data": data}
201
  return jsonify(response)
202
 
203
+ def apikey_required(func):
204
+ @wraps(func)
205
+ def decorated_function(*args, **kwargs):
206
+ token = flask_request.headers.get('Authorization').split()[1]
207
+ objs = APIToken.query(token=token)
208
+ if not objs:
209
+ return build_error_result(
210
+ error_msg='API-KEY is invalid!', retcode=RetCode.FORBIDDEN
211
+ )
212
+ kwargs['tenant_id'] = objs[0].tenant_id
213
+ return func(*args, **kwargs)
214
+
215
+ return decorated_function
216
+
217
+
218
+ def build_error_result(retcode=RetCode.FORBIDDEN, error_msg='success'):
219
+ response = {"error_code": retcode, "error_msg": error_msg}
220
+ response = jsonify(response)
221
+ response.status_code = retcode
222
+ return response
223
+
224
 
225
  def construct_response(retcode=RetCode.SUCCESS,
226
  retmsg='success', data=None, auth=None):