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 +62 -0
- api/settings.py +2 -0
- api/utils/api_utils.py +21 -0
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):
|