pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Unified API response json schema (#3170)

Browse Source 
### What problem does this PR solve?

Unified API response json schema

### Type of change

- [x] Refactoring
This commit is contained in:
Zhichang Yu
2024-11-05 11:02:31 +08:00
committed by GitHub
parent 339639a9db
commit 185c6a0c71
53 changed files with 1458 additions and 1470 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
api/apps/api_app.py
View File

@ -52,7 +52,7 @@ def new_token():
try:
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
tenant_id = tenants[0].tenant_id
obj = {"tenant_id": tenant_id, "token": generate_confirmation_token(tenant_id),
@ -68,7 +68,7 @@ def new_token():
obj["dialog_id"] = req["dialog_id"]
if not APITokenService.save(**obj):
return get_data_error_result(retmsg="Fail to new a dialog!")
return get_data_error_result(message="Fail to new a dialog!")
return get_json_result(data=obj)
except Exception as e:
@ -81,7 +81,7 @@ def token_list():
try:
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
id = request.args["dialog_id"] if "dialog_id" in request.args else request.args["canvas_id"]
objs = APITokenService.query(tenant_id=tenants[0].tenant_id, dialog_id=id)
@ -110,7 +110,7 @@ def stats():
try:
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
objs = API4ConversationService.stats(
tenants[0].tenant_id,
request.args.get(
@ -141,7 +141,7 @@ def set_conversation():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
req = request.json
try:
if objs[0].source == "agent":
@ -163,7 +163,7 @@ def set_conversation():
else:
e, dia = DialogService.get_by_id(objs[0].dialog_id)
if not e:
return get_data_error_result(retmsg="Dialog not found")
return get_data_error_result(message="Dialog not found")
conv = {
"id": get_uuid(),
"dialog_id": dia.id,
@ -183,11 +183,11 @@ def completion():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
req = request.json
e, conv = API4ConversationService.get_by_id(req["conversation_id"])
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
if "quote" not in req: req["quote"] = False
msg = []
@ -257,7 +257,7 @@ def completion():
ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
fillin_conv(ans)
rename_field(ans)
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans},
ensure_ascii=False) + "\n\n"
canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
@ -267,10 +267,10 @@ def completion():
cvs.dsl = json.loads(str(canvas))
API4ConversationService.append_message(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(sse(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
@ -295,7 +295,7 @@ def completion():
conv.message.append(msg[-1])
e, dia = DialogService.get_by_id(conv.dialog_id)
if not e:
return get_data_error_result(retmsg="Dialog not found!")
return get_data_error_result(message="Dialog not found!")
del req["conversation_id"]
del req["messages"]
@ -310,14 +310,14 @@ def completion():
for ans in chat(dia, msg, True, **req):
fillin_conv(ans)
rename_field(ans)
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans},
ensure_ascii=False) + "\n\n"
API4ConversationService.append_message(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
if req.get("stream", True):
resp = Response(stream(), mimetype="text/event-stream")
@ -347,17 +347,17 @@ def get(conversation_id):
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
try:
e, conv = API4ConversationService.get_by_id(conversation_id)
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
conv = conv.to_dict()
if token != APIToken.query(dialog_id=conv['dialog_id'])[0].token:
return get_json_result(data=False, retmsg='Token is not valid for this conversation_id!"',
retcode=RetCode.AUTHENTICATION_ERROR)
return get_json_result(data=False, message='Token is not valid for this conversation_id!"',
code=RetCode.AUTHENTICATION_ERROR)
for referenct_i in conv['reference']:
if referenct_i is None or len(referenct_i) == 0:
@ -378,7 +378,7 @@ def upload():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
kb_name = request.form.get("kb_name").strip()
tenant_id = objs[0].tenant_id
@ -387,19 +387,19 @@ def upload():
e, kb = KnowledgebaseService.get_by_name(kb_name, tenant_id)
if not e:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
kb_id = kb.id
except Exception as e:
return server_error_response(e)
if 'file' not in request.files:
return get_json_result(
data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file part!', code=RetCode.ARGUMENT_ERROR)
file = request.files['file']
if file.filename == '':
return get_json_result(
data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file selected!', code=RetCode.ARGUMENT_ERROR)
root_folder = FileService.get_root_folder(tenant_id)
pf_id = root_folder["id"]
@ -410,7 +410,7 @@ def upload():
try:
if DocumentService.get_doc_count(kb.tenant_id) >= int(os.environ.get('MAX_FILE_NUM_PER_USER', 8192)):
return get_data_error_result(
retmsg="Exceed the maximum file number of a free user!")
message="Exceed the maximum file number of a free user!")
filename = duplicate_name(
DocumentService.query,
@ -419,7 +419,7 @@ def upload():
filetype = filename_type(filename)
if not filetype:
return get_data_error_result(
retmsg="This type of file has not been supported yet!")
message="This type of file has not been supported yet!")
location = filename
while STORAGE_IMPL.obj_exist(kb_id, location):
@ -468,7 +468,7 @@ def upload():
# if str(req["run"]) == TaskStatus.CANCEL.value:
tenant_id = DocumentService.get_tenant_id(doc["id"])
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
# e, doc = DocumentService.get_by_id(doc["id"])
TaskService.filter_delete([Task.doc_id == doc["id"]])
@ -490,17 +490,17 @@ def upload_parse():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
if 'file' not in request.files:
return get_json_result(
data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file part!', code=RetCode.ARGUMENT_ERROR)
file_objs = request.files.getlist('file')
for file_obj in file_objs:
if file_obj.filename == '':
return get_json_result(
data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file selected!', code=RetCode.ARGUMENT_ERROR)
doc_ids = doc_upload_and_parse(request.form.get("conversation_id"), file_objs, objs[0].tenant_id)
return get_json_result(data=doc_ids)
@ -513,7 +513,7 @@ def list_chunks():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
req = request.json
@ -527,7 +527,7 @@ def list_chunks():
doc_id = req['doc_id']
else:
return get_json_result(
data=False, retmsg="Can't find doc_name or doc_id"
data=False, message="Can't find doc_name or doc_id"
)
res = retrievaler.chunk_list(doc_id=doc_id, tenant_id=tenant_id)
@ -552,7 +552,7 @@ def list_kb_docs():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
req = request.json
tenant_id = objs[0].tenant_id
@ -562,7 +562,7 @@ def list_kb_docs():
e, kb = KnowledgebaseService.get_by_name(kb_name, tenant_id)
if not e:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
kb_id = kb.id
except Exception as e:
@ -591,7 +591,7 @@ def docinfos():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
req = request.json
doc_ids = req["doc_ids"]
docs = DocumentService.get_by_ids(doc_ids)
@ -605,7 +605,7 @@ def document_rm():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
tenant_id = objs[0].tenant_id
req = request.json
@ -617,7 +617,7 @@ def document_rm():
if not doc_ids:
return get_json_result(
data=False, retmsg="Can't find doc_names or doc_ids"
data=False, message="Can't find doc_names or doc_ids"
)
except Exception as e:
@ -632,16 +632,16 @@ def document_rm():
try:
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(doc_id)
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
b, n = File2DocumentService.get_storage_address(doc_id=doc_id)
if not DocumentService.remove_document(doc, tenant_id):
return get_data_error_result(
retmsg="Database error (Document removal)!")
message="Database error (Document removal)!")
f2d = File2DocumentService.get_by_document_id(doc_id)
FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
@ -652,7 +652,7 @@ def document_rm():
errors += str(e)
if errors:
return get_json_result(data=False, retmsg=errors, retcode=RetCode.SERVER_ERROR)
return get_json_result(data=False, message=errors, code=RetCode.SERVER_ERROR)
return get_json_result(data=True)
@ -667,11 +667,11 @@ def completion_faq():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
e, conv = API4ConversationService.get_by_id(req["conversation_id"])
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
if "quote" not in req: req["quote"] = True
msg = []
@ -752,7 +752,7 @@ def completion_faq():
conv.message.append(msg[-1])
e, dia = DialogService.get_by_id(conv.dialog_id)
if not e:
return get_data_error_result(retmsg="Dialog not found!")
return get_data_error_result(message="Dialog not found!")
del req["conversation_id"]
if not conv.reference:
@ -804,7 +804,7 @@ def retrieval():
objs = APIToken.query(token=token)
if not objs:
return get_json_result(
data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Token is not valid!"', code=RetCode.AUTHENTICATION_ERROR)
req = request.json
kb_ids = req.get("kb_id",[])
@ -821,7 +821,7 @@ def retrieval():
embd_nms = list(set([kb.embd_id for kb in kbs]))
if len(embd_nms) != 1:
return get_json_result(
data=False, retmsg='Knowledge bases use different embedding models or does not exist."', retcode=RetCode.AUTHENTICATION_ERROR)
data=False, message='Knowledge bases use different embedding models or does not exist."', code=RetCode.AUTHENTICATION_ERROR)
embd_mdl = TenantLLMService.model_instance(
kbs[0].tenant_id, LLMType.EMBEDDING.value, llm_name=kbs[0].embd_id)
@ -841,6 +841,6 @@ def retrieval():
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'No chunk found! Check the chunk status please!',
retcode=RetCode.DATA_ERROR)
return get_json_result(data=False, message='No chunk found! Check the chunk status please!',
code=RetCode.DATA_ERROR)
return server_error_response(e)

30
api/apps/canvas_app.py
View File

@ -46,8 +46,8 @@ def rm():
for i in request.json["canvas_ids"]:
if not UserCanvasService.query(user_id=current_user.id,id=i):
return get_json_result(
data=False, retmsg=f'Only owner of canvas authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
UserCanvasService.delete_by_id(i)
return get_json_result(data=True)
@ -66,12 +66,12 @@ def save():
return server_error_response(ValueError("Duplicated title."))
req["id"] = get_uuid()
if not UserCanvasService.save(**req):
return get_data_error_result(retmsg="Fail to save canvas.")
return get_data_error_result(message="Fail to save canvas.")
else:
if not UserCanvasService.query(user_id=current_user.id, id=req["id"]):
return get_json_result(
data=False, retmsg=f'Only owner of canvas authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
UserCanvasService.update_by_id(req["id"], req)
return get_json_result(data=req)
@ -81,7 +81,7 @@ def save():
def get(canvas_id):
e, c = UserCanvasService.get_by_id(canvas_id)
if not e:
return get_data_error_result(retmsg="canvas not found.")
return get_data_error_result(message="canvas not found.")
return get_json_result(data=c.to_dict())
@ -93,11 +93,11 @@ def run():
stream = req.get("stream", True)
e, cvs = UserCanvasService.get_by_id(req["id"])
if not e:
return get_data_error_result(retmsg="canvas not found.")
return get_data_error_result(message="canvas not found.")
if not UserCanvasService.query(user_id=current_user.id, id=req["id"]):
return get_json_result(
data=False, retmsg=f'Only owner of canvas authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
if not isinstance(cvs.dsl, str):
cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
@ -130,7 +130,7 @@ def run():
for k in ans.keys():
final_ans[k] = ans[k]
ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
canvas.history.append(("assistant", final_ans["content"]))
@ -139,10 +139,10 @@ def run():
cvs.dsl = json.loads(str(canvas))
UserCanvasService.update_by_id(req["id"], cvs.to_dict())
except Exception as e:
yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(sse(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
@ -168,11 +168,11 @@ def reset():
try:
e, user_canvas = UserCanvasService.get_by_id(req["id"])
if not e:
return get_data_error_result(retmsg="canvas not found.")
return get_data_error_result(message="canvas not found.")
if not UserCanvasService.query(user_id=current_user.id, id=req["id"]):
return get_json_result(
data=False, retmsg=f'Only owner of canvas authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
canvas = Canvas(json.dumps(user_canvas.dsl), current_user.id)
canvas.reset()

45
api/apps/chunk_app.py
View File

@ -15,7 +15,6 @@
#
import datetime
import json
import traceback
from flask import request
from flask_login import login_required, current_user
@ -50,10 +49,10 @@ def list_chunk():
try:
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
query = {
"doc_ids": [doc_id], "page": page, "size": size, "question": question, "sort": True
}
@ -84,8 +83,8 @@ def list_chunk():
return get_json_result(data=res)
except Exception as e:
if str(e).find("not_found") > 0:
return get_json_result(data=False, retmsg=f'No chunk found!',
retcode=RetCode.DATA_ERROR)
return get_json_result(data=False, message='No chunk found!',
code=RetCode.DATA_ERROR)
return server_error_response(e)
@ -96,7 +95,7 @@ def get():
try:
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
res = ELASTICSEARCH.get(
chunk_id, search.index_name(
tenants[0].tenant_id))
@ -115,8 +114,8 @@ def get():
return get_json_result(data=res)
except Exception as e:
if str(e).find("NotFoundError") >= 0:
return get_json_result(data=False, retmsg=f'Chunk not found!',
retcode=RetCode.DATA_ERROR)
return get_json_result(data=False, message='Chunk not found!',
code=RetCode.DATA_ERROR)
return server_error_response(e)
@ -139,14 +138,14 @@ def set():
try:
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
embd_id = DocumentService.get_embd_id(req["doc_id"])
embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embd_id)
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
if doc.parser_id == ParserType.QA:
arr = [
@ -155,7 +154,7 @@ def set():
req["content_with_weight"]) if len(t) > 1]
if len(arr) != 2:
return get_data_error_result(
retmsg="Q&A must be separated by TAB/ENTER key.")
message="Q&A must be separated by TAB/ENTER key.")
q, a = rmPrefix(arr[0]), rmPrefix(arr[1])
d = beAdoc(d, arr[0], arr[1], not any(
[rag_tokenizer.is_chinese(t) for t in q + a]))
@ -177,10 +176,10 @@ def switch():
try:
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="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)):
return get_data_error_result(retmsg="Index updating failure")
return get_data_error_result(message="Index updating failure")
return get_json_result(data=True)
except Exception as e:
return server_error_response(e)
@ -194,10 +193,10 @@ def rm():
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_data_error_result(message="Index updating failure")
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
deleted_chunk_ids = req["chunk_ids"]
chunk_number = len(deleted_chunk_ids)
DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
@ -225,14 +224,14 @@ def create():
try:
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="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!")
return get_data_error_result(message="Tenant not found!")
embd_id = DocumentService.get_embd_id(req["doc_id"])
embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING.value, embd_id)
@ -273,12 +272,12 @@ def retrieval_test():
break
else:
return get_json_result(
data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of knowledgebase authorized for this operation.',
code=RetCode.OPERATING_ERROR)
e, kb = KnowledgebaseService.get_by_id(kb_id[0])
if not e:
return get_data_error_result(retmsg="Knowledgebase not found!")
return get_data_error_result(message="Knowledgebase not found!")
embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
@ -301,8 +300,8 @@ def retrieval_test():
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'No chunk found! Check the chunk status please!',
retcode=RetCode.DATA_ERROR)
return get_json_result(data=False, message='No chunk found! Check the chunk status please!',
code=RetCode.DATA_ERROR)
return server_error_response(e)
@ -321,7 +320,7 @@ def knowledge_graph():
ty = sres.field[id]["knowledge_graph_kwd"]
try:
content_json = json.loads(sres.field[id]["content_with_weight"])
except Exception as e:
except Exception:
continue
if ty == 'mind_map':

52
api/apps/conversation_app.py
View File

@ -42,11 +42,11 @@ def set_conversation():
del req["conversation_id"]
try:
if not ConversationService.update_by_id(conv_id, req):
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="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!")
message="Fail to update a conversation!")
conv = conv.to_dict()
return get_json_result(data=conv)
except Exception as e:
@ -55,7 +55,7 @@ def set_conversation():
try:
e, dia = DialogService.get_by_id(req["dialog_id"])
if not e:
return get_data_error_result(retmsg="Dialog not found")
return get_data_error_result(message="Dialog not found")
conv = {
"id": conv_id,
"dialog_id": req["dialog_id"],
@ -65,7 +65,7 @@ def set_conversation():
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!")
return get_data_error_result(message="Fail to new a conversation!")
conv = conv.to_dict()
return get_json_result(data=conv)
except Exception as e:
@ -79,15 +79,15 @@ def get():
try:
e, conv = ConversationService.get_by_id(conv_id)
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
tenants = UserTenantService.query(user_id=current_user.id)
for tenant in tenants:
if DialogService.query(tenant_id=tenant.tenant_id, id=conv.dialog_id):
break
else:
return get_json_result(
data=False, retmsg=f'Only owner of conversation authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of conversation authorized for this operation.',
code=RetCode.OPERATING_ERROR)
conv = conv.to_dict()
return get_json_result(data=conv)
except Exception as e:
@ -102,15 +102,15 @@ def rm():
for cid in conv_ids:
exist, conv = ConversationService.get_by_id(cid)
if not exist:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
tenants = UserTenantService.query(user_id=current_user.id)
for tenant in tenants:
if DialogService.query(tenant_id=tenant.tenant_id, id=conv.dialog_id):
break
else:
return get_json_result(
data=False, retmsg=f'Only owner of conversation authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of conversation authorized for this operation.',
code=RetCode.OPERATING_ERROR)
ConversationService.delete_by_id(cid)
return get_json_result(data=True)
except Exception as e:
@ -124,8 +124,8 @@ def list_convsersation():
try:
if not DialogService.query(tenant_id=current_user.id, id=dialog_id):
return get_json_result(
data=False, retmsg=f'Only owner of dialog authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of dialog authorized for this operation.',
code=RetCode.OPERATING_ERROR)
convs = ConversationService.query(
dialog_id=dialog_id,
order_by=ConversationService.model.create_time,
@ -152,11 +152,11 @@ def completion():
try:
e, conv = ConversationService.get_by_id(req["conversation_id"])
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
conv.message = deepcopy(req["messages"])
e, dia = DialogService.get_by_id(conv.dialog_id)
if not e:
return get_data_error_result(retmsg="Dialog not found!")
return get_data_error_result(message="Dialog not found!")
del req["conversation_id"]
del req["messages"]
@ -180,14 +180,14 @@ def completion():
try:
for ans in chat(dia, msg, True, **req):
fillin_conv(ans)
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
ConversationService.update_by_id(conv.id, conv.to_dict())
except Exception as e:
traceback.print_exc()
yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
if req.get("stream", True):
resp = Response(stream(), mimetype="text/event-stream")
@ -217,11 +217,11 @@ def tts():
tenants = TenantService.get_info_by(current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
tts_id = tenants[0]["tts_id"]
if not tts_id:
return get_data_error_result(retmsg="No default TTS model is set")
return get_data_error_result(message="No default TTS model is set")
tts_mdl = LLMBundle(tenants[0]["tenant_id"], LLMType.TTS, tts_id)
@ -231,7 +231,7 @@ def tts():
for chunk in tts_mdl.tts(txt):
yield chunk
except Exception as e:
yield ("data:" + json.dumps({"retcode": 500, "retmsg": str(e),
yield ("data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e)}},
ensure_ascii=False)).encode('utf-8')
@ -250,7 +250,7 @@ def delete_msg():
req = request.json
e, conv = ConversationService.get_by_id(req["conversation_id"])
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
conv = conv.to_dict()
for i, msg in enumerate(conv["message"]):
@ -273,7 +273,7 @@ def thumbup():
req = request.json
e, conv = ConversationService.get_by_id(req["conversation_id"])
if not e:
return get_data_error_result(retmsg="Conversation not found!")
return get_data_error_result(message="Conversation not found!")
up_down = req.get("set")
feedback = req.get("feedback", "")
conv = conv.to_dict()
@ -301,12 +301,12 @@ def ask_about():
nonlocal req, uid
try:
for ans in ask(req["question"], req["kb_ids"], uid):
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
except Exception as e:
yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(stream(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
@ -324,7 +324,7 @@ def mindmap():
kb_ids = req["kb_ids"]
e, kb = KnowledgebaseService.get_by_id(kb_ids[0])
if not e:
return get_data_error_result(retmsg="Knowledgebase not found!")
return get_data_error_result(message="Knowledgebase not found!")
embd_mdl = TenantLLMService.model_instance(
kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)

20
api/apps/dialog_app.py
View File

@ -68,17 +68,17 @@ def set_dialog():
continue
if prompt_config["system"].find("{%s}" % p["key"]) < 0:
return get_data_error_result(
retmsg="Parameter '{}' is not used".format(p["key"]))
message="Parameter '{}' is not used".format(p["key"]))
try:
e, tenant = TenantService.get_by_id(current_user.id)
if not e:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
llm_id = req.get("llm_id", tenant.llm_id)
if not dialog_id:
if not req.get("kb_ids"):
return get_data_error_result(
retmsg="Fail! Please select knowledgebase!")
message="Fail! Please select knowledgebase!")
dia = {
"id": get_uuid(),
"tenant_id": current_user.id,
@ -96,20 +96,20 @@ def set_dialog():
"icon": icon
}
if not DialogService.save(**dia):
return get_data_error_result(retmsg="Fail to new a dialog!")
return get_data_error_result(message="Fail to new a dialog!")
e, dia = DialogService.get_by_id(dia["id"])
if not e:
return get_data_error_result(retmsg="Fail to new a dialog!")
return get_data_error_result(message="Fail to new a dialog!")
return get_json_result(data=dia.to_json())
else:
del req["dialog_id"]
if "kb_names" in req:
del req["kb_names"]
if not DialogService.update_by_id(dialog_id, req):
return get_data_error_result(retmsg="Dialog not found!")
return get_data_error_result(message="Dialog not found!")
e, dia = DialogService.get_by_id(dialog_id)
if not e:
return get_data_error_result(retmsg="Fail to update a dialog!")
return get_data_error_result(message="Fail to update a dialog!")
dia = dia.to_dict()
dia["kb_ids"], dia["kb_names"] = get_kb_names(dia["kb_ids"])
return get_json_result(data=dia)
@ -124,7 +124,7 @@ def get():
try:
e, dia = DialogService.get_by_id(dialog_id)
if not e:
return get_data_error_result(retmsg="Dialog not found!")
return get_data_error_result(message="Dialog not found!")
dia = dia.to_dict()
dia["kb_ids"], dia["kb_names"] = get_kb_names(dia["kb_ids"])
return get_json_result(data=dia)
@ -174,8 +174,8 @@ def rm():
break
else:
return get_json_result(
data=False, retmsg=f'Only owner of dialog authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of dialog authorized for this operation.',
code=RetCode.OPERATING_ERROR)
dialog_list.append({"id": id,"status":StatusEnum.INVALID.value})
DialogService.update_many_by_id(dialog_list)
return get_json_result(data=True)

96
api/apps/document_app.py
View File

@ -49,16 +49,16 @@ def upload():
kb_id = request.form.get("kb_id")
if not kb_id:
return get_json_result(
data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='Lack of "KB ID"', code=RetCode.ARGUMENT_ERROR)
if 'file' not in request.files:
return get_json_result(
data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file part!', code=RetCode.ARGUMENT_ERROR)
file_objs = request.files.getlist('file')
for file_obj in file_objs:
if file_obj.filename == '':
return get_json_result(
data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file selected!', code=RetCode.ARGUMENT_ERROR)
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
@ -67,7 +67,7 @@ def upload():
err, _ = FileService.upload_document(kb, file_objs, current_user.id)
if err:
return get_json_result(
data=False, retmsg="\n".join(err), retcode=RetCode.SERVER_ERROR)
data=False, message="\n".join(err), code=RetCode.SERVER_ERROR)
return get_json_result(data=True)
@ -78,12 +78,12 @@ def web_crawl():
kb_id = request.form.get("kb_id")
if not kb_id:
return get_json_result(
data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='Lack of "KB ID"', code=RetCode.ARGUMENT_ERROR)
name = request.form.get("name")
url = request.form.get("url")
if not is_valid_url(url):
return get_json_result(
data=False, retmsg='The URL format is invalid', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='The URL format is invalid', code=RetCode.ARGUMENT_ERROR)
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
raise LookupError("Can't find this knowledgebase!")
@ -145,17 +145,17 @@ def create():
kb_id = req["kb_id"]
if not kb_id:
return get_json_result(
data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='Lack of "KB ID"', code=RetCode.ARGUMENT_ERROR)
try:
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
if DocumentService.query(name=req["name"], kb_id=kb_id):
return get_data_error_result(
retmsg="Duplicated document name in the same knowledgebase.")
message="Duplicated document name in the same knowledgebase.")
doc = DocumentService.insert({
"id": get_uuid(),
@ -179,7 +179,7 @@ def list_docs():
kb_id = request.args.get("kb_id")
if not kb_id:
return get_json_result(
data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='Lack of "KB ID"', code=RetCode.ARGUMENT_ERROR)
tenants = UserTenantService.query(user_id=current_user.id)
for tenant in tenants:
if KnowledgebaseService.query(
@ -187,8 +187,8 @@ def list_docs():
break
else:
return get_json_result(
data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of knowledgebase authorized for this operation.',
code=RetCode.OPERATING_ERROR)
keywords = request.args.get("keywords", "")
page_number = int(request.args.get("page", 1))
@ -217,8 +217,8 @@ def docinfos():
if not DocumentService.accessible(doc_id, current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
docs = DocumentService.get_by_ids(doc_ids)
return get_json_result(data=list(docs.dicts()))
@ -230,7 +230,7 @@ def thumbnails():
doc_ids = request.args.get("doc_ids").split(",")
if not doc_ids:
return get_json_result(
data=False, retmsg='Lack of "Document ID"', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='Lack of "Document ID"', code=RetCode.ARGUMENT_ERROR)
try:
docs = DocumentService.get_thumbnails(doc_ids)
@ -252,28 +252,28 @@ def change_status():
if str(req["status"]) not in ["0", "1"]:
return get_json_result(
data=False,
retmsg='"Status" must be either 0 or 1!',
retcode=RetCode.ARGUMENT_ERROR)
message='"Status" must be either 0 or 1!',
code=RetCode.ARGUMENT_ERROR)
if not DocumentService.accessible(req["doc_id"], current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR)
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR)
try:
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
e, kb = KnowledgebaseService.get_by_id(doc.kb_id)
if not e:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
if not DocumentService.update_by_id(
req["doc_id"], {"status": str(req["status"])}):
return get_data_error_result(
retmsg="Database error (Document update)!")
message="Database error (Document update)!")
if str(req["status"]) == "0":
ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=req["doc_id"]),
@ -304,8 +304,8 @@ def rm():
if not DocumentService.accessible4deletion(doc_id, current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
root_folder = FileService.get_root_folder(current_user.id)
@ -316,16 +316,16 @@ def rm():
try:
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(doc_id)
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
b, n = File2DocumentService.get_storage_address(doc_id=doc_id)
if not DocumentService.remove_document(doc, tenant_id):
return get_data_error_result(
retmsg="Database error (Document removal)!")
message="Database error (Document removal)!")
f2d = File2DocumentService.get_by_document_id(doc_id)
FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
@ -336,7 +336,7 @@ def rm():
errors += str(e)
if errors:
return get_json_result(data=False, retmsg=errors, retcode=RetCode.SERVER_ERROR)
return get_json_result(data=False, message=errors, code=RetCode.SERVER_ERROR)
return get_json_result(data=True)
@ -350,8 +350,8 @@ def run():
if not DocumentService.accessible(doc_id, current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
try:
for id in req["doc_ids"]:
@ -364,7 +364,7 @@ def run():
# if str(req["run"]) == TaskStatus.CANCEL.value:
tenant_id = DocumentService.get_tenant_id(id)
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
ELASTICSEARCH.deleteByQuery(
Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
@ -389,28 +389,28 @@ def rename():
if not DocumentService.accessible(req["doc_id"], current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
try:
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
doc.name.lower()).suffix:
return get_json_result(
data=False,
retmsg="The extension of file can't be changed",
retcode=RetCode.ARGUMENT_ERROR)
message="The extension of file can't be changed",
code=RetCode.ARGUMENT_ERROR)
for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
if d.name == req["name"]:
return get_data_error_result(
retmsg="Duplicated document name in the same knowledgebase.")
message="Duplicated document name in the same knowledgebase.")
if not DocumentService.update_by_id(
req["doc_id"], {"name": req["name"]}):
return get_data_error_result(
retmsg="Database error (Document rename)!")
message="Database error (Document rename)!")
informs = File2DocumentService.get_by_document_id(req["doc_id"])
if informs:
@ -428,7 +428,7 @@ def get(doc_id):
try:
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
b, n = File2DocumentService.get_storage_address(doc_id=doc_id)
response = flask.make_response(STORAGE_IMPL.get(b, n))
@ -456,13 +456,13 @@ def change_parser():
if not DocumentService.accessible(req["doc_id"], current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
try:
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
if doc.parser_id.lower() == req["parser_id"].lower():
if "parser_config" in req:
if req["parser_config"] == doc.parser_config:
@ -473,23 +473,23 @@ def change_parser():
if ((doc.type == FileType.VISUAL and req["parser_id"] != "picture")
or (re.search(
r"\.(ppt|pptx|pages)$", doc.name) and req["parser_id"] != "presentation")):
return get_data_error_result(retmsg="Not supported yet!")
return get_data_error_result(message="Not supported yet!")
e = DocumentService.update_by_id(doc.id,
{"parser_id": req["parser_id"], "progress": 0, "progress_msg": "",
"run": TaskStatus.UNSTART.value})
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
if "parser_config" in req:
DocumentService.update_parser_config(doc.id, req["parser_config"])
if doc.token_num > 0:
e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num * -1, doc.chunk_num * -1,
doc.process_duation * -1)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
ELASTICSEARCH.deleteByQuery(
Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id))
@ -516,13 +516,13 @@ def get_image(image_id):
def upload_and_parse():
if 'file' not in request.files:
return get_json_result(
data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file part!', code=RetCode.ARGUMENT_ERROR)
file_objs = request.files.getlist('file')
for file_obj in file_objs:
if file_obj.filename == '':
return get_json_result(
data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file selected!', code=RetCode.ARGUMENT_ERROR)
doc_ids = doc_upload_and_parse(request.form.get("conversation_id"), file_objs, current_user.id)

26
api/apps/file2document_app.py
View File

@ -13,9 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License
#
from elasticsearch_dsl import Q
from api.db.db_models import File2Document
from api.db.services.file2document_service import File2DocumentService
from api.db.services.file_service import FileService
@ -28,8 +26,6 @@ from api.db import FileType
from api.db.services.document_service import DocumentService
from api.settings import RetCode
from api.utils.api_utils import get_json_result
from rag.nlp import search
from rag.utils.es_conn import ELASTICSEARCH
@manager.route('/convert', methods=['POST'])
@ -54,13 +50,13 @@ def convert():
doc_id = inform.document_id
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(doc_id)
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
if not DocumentService.remove_document(doc, tenant_id):
return get_data_error_result(
retmsg="Database error (Document removal)!")
message="Database error (Document removal)!")
File2DocumentService.delete_by_file_id(id)
# insert
@ -68,11 +64,11 @@ def convert():
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
e, file = FileService.get_by_id(id)
if not e:
return get_data_error_result(
retmsg="Can't find this file!")
message="Can't find this file!")
doc = DocumentService.insert({
"id": get_uuid(),
@ -104,26 +100,26 @@ def rm():
file_ids = req["file_ids"]
if not file_ids:
return get_json_result(
data=False, retmsg='Lack of "Files ID"', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='Lack of "Files ID"', code=RetCode.ARGUMENT_ERROR)
try:
for file_id in file_ids:
informs = File2DocumentService.get_by_file_id(file_id)
if not informs:
return get_data_error_result(retmsg="Inform not found!")
return get_data_error_result(message="Inform not found!")
for inform in informs:
if not inform:
return get_data_error_result(retmsg="Inform not found!")
return get_data_error_result(message="Inform not found!")
File2DocumentService.delete_by_file_id(file_id)
doc_id = inform.document_id
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(doc_id)
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
if not DocumentService.remove_document(doc, tenant_id):
return get_data_error_result(
retmsg="Database error (Document removal)!")
message="Database error (Document removal)!")
return get_json_result(data=True)
except Exception as e:
return server_error_response(e)

61
api/apps/file_app.py
View File

@ -18,7 +18,6 @@ import pathlib
import re
import flask
from elasticsearch_dsl import Q
from flask import request
from flask_login import login_required, current_user
@ -32,8 +31,6 @@ from api.db.services.file_service import FileService
from api.settings import RetCode
from api.utils.api_utils import get_json_result
from api.utils.file_utils import filename_type
from rag.nlp import search
from rag.utils.es_conn import ELASTICSEARCH
from rag.utils.storage_factory import STORAGE_IMPL
@ -49,24 +46,24 @@ def upload():
if 'file' not in request.files:
return get_json_result(
data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file part!', code=RetCode.ARGUMENT_ERROR)
file_objs = request.files.getlist('file')
for file_obj in file_objs:
if file_obj.filename == '':
return get_json_result(
data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
data=False, message='No file selected!', code=RetCode.ARGUMENT_ERROR)
file_res = []
try:
for file_obj in file_objs:
e, file = FileService.get_by_id(pf_id)
if not e:
return get_data_error_result(
retmsg="Can't find this folder!")
message="Can't find this folder!")
MAX_FILE_NUM_PER_USER = int(os.environ.get('MAX_FILE_NUM_PER_USER', 0))
if MAX_FILE_NUM_PER_USER > 0 and DocumentService.get_doc_count(current_user.id) >= MAX_FILE_NUM_PER_USER:
return get_data_error_result(
retmsg="Exceed the maximum file number of a free user!")
message="Exceed the maximum file number of a free user!")
# split file name path
if not file_obj.filename:
@ -85,13 +82,13 @@ def upload():
if file_len != len_id_list:
e, file = FileService.get_by_id(file_id_list[len_id_list - 1])
if not e:
return get_data_error_result(retmsg="Folder not found!")
return get_data_error_result(message="Folder not found!")
last_folder = FileService.create_folder(file, file_id_list[len_id_list - 1], file_obj_names,
len_id_list)
else:
e, file = FileService.get_by_id(file_id_list[len_id_list - 2])
if not e:
return get_data_error_result(retmsg="Folder not found!")
return get_data_error_result(message="Folder not found!")
last_folder = FileService.create_folder(file, file_id_list[len_id_list - 2], file_obj_names,
len_id_list)
@ -137,10 +134,10 @@ def create():
try:
if not FileService.is_parent_folder_exist(pf_id):
return get_json_result(
data=False, retmsg="Parent Folder Doesn't Exist!", retcode=RetCode.OPERATING_ERROR)
data=False, message="Parent Folder Doesn't Exist!", code=RetCode.OPERATING_ERROR)
if FileService.query(name=req["name"], parent_id=pf_id):
return get_data_error_result(
retmsg="Duplicated folder name in the same folder.")
message="Duplicated folder name in the same folder.")
if input_file_type == FileType.FOLDER.value:
file_type = FileType.FOLDER.value
@ -181,14 +178,14 @@ def list_files():
try:
e, file = FileService.get_by_id(pf_id)
if not e:
return get_data_error_result(retmsg="Folder not found!")
return get_data_error_result(message="Folder not found!")
files, total = FileService.get_by_pf_id(
current_user.id, pf_id, page_number, items_per_page, orderby, desc, keywords)
parent_folder = FileService.get_parent_folder(pf_id)
if not FileService.get_parent_folder(pf_id):
return get_json_result(retmsg="File not found!")
return get_json_result(message="File not found!")
return get_json_result(data={"total": total, "files": files, "parent_folder": parent_folder.to_json()})
except Exception as e:
@ -212,7 +209,7 @@ def get_parent_folder():
try:
e, file = FileService.get_by_id(file_id)
if not e:
return get_data_error_result(retmsg="Folder not found!")
return get_data_error_result(message="Folder not found!")
parent_folder = FileService.get_parent_folder(file_id)
return get_json_result(data={"parent_folder": parent_folder.to_json()})
@ -227,7 +224,7 @@ def get_all_parent_folders():
try:
e, file = FileService.get_by_id(file_id)
if not e:
return get_data_error_result(retmsg="Folder not found!")
return get_data_error_result(message="Folder not found!")
parent_folders = FileService.get_all_parent_folders(file_id)
parent_folders_res = []
@ -248,9 +245,9 @@ def rm():
for file_id in file_ids:
e, file = FileService.get_by_id(file_id)
if not e:
return get_data_error_result(retmsg="File or Folder not found!")
return get_data_error_result(message="File or Folder not found!")
if not file.tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
if file.source_type == FileSource.KNOWLEDGEBASE:
continue
@ -259,13 +256,13 @@ def rm():
for inner_file_id in file_id_list:
e, file = FileService.get_by_id(inner_file_id)
if not e:
return get_data_error_result(retmsg="File not found!")
return get_data_error_result(message="File not found!")
STORAGE_IMPL.rm(file.parent_id, file.location)
FileService.delete_folder_by_pf_id(current_user.id, file_id)
else:
if not FileService.delete(file):
return get_data_error_result(
retmsg="Database error (File removal)!")
message="Database error (File removal)!")
# delete file2document
informs = File2DocumentService.get_by_file_id(file_id)
@ -273,13 +270,13 @@ def rm():
doc_id = inform.document_id
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(doc_id)
if not tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
if not DocumentService.remove_document(doc, tenant_id):
return get_data_error_result(
retmsg="Database error (Document removal)!")
message="Database error (Document removal)!")
File2DocumentService.delete_by_file_id(file_id)
return get_json_result(data=True)
@ -295,30 +292,30 @@ def rename():
try:
e, file = FileService.get_by_id(req["file_id"])
if not e:
return get_data_error_result(retmsg="File not found!")
return get_data_error_result(message="File not found!")
if file.type != FileType.FOLDER.value \
and pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
file.name.lower()).suffix:
return get_json_result(
data=False,
retmsg="The extension of file can't be changed",
retcode=RetCode.ARGUMENT_ERROR)
message="The extension of file can't be changed",
code=RetCode.ARGUMENT_ERROR)
for file in FileService.query(name=req["name"], pf_id=file.parent_id):
if file.name == req["name"]:
return get_data_error_result(
retmsg="Duplicated file name in the same folder.")
message="Duplicated file name in the same folder.")
if not FileService.update_by_id(
req["file_id"], {"name": req["name"]}):
return get_data_error_result(
retmsg="Database error (File rename)!")
message="Database error (File rename)!")
informs = File2DocumentService.get_by_file_id(req["file_id"])
if informs:
if not DocumentService.update_by_id(
informs[0].document_id, {"name": req["name"]}):
return get_data_error_result(
retmsg="Database error (Document rename)!")
message="Database error (Document rename)!")
return get_json_result(data=True)
except Exception as e:
@ -331,7 +328,7 @@ def get(file_id):
try:
e, file = FileService.get_by_id(file_id)
if not e:
return get_data_error_result(retmsg="Document not found!")
return get_data_error_result(message="Document not found!")
b, n = File2DocumentService.get_storage_address(file_id=file_id)
response = flask.make_response(STORAGE_IMPL.get(b, n))
ext = re.search(r"\.([^.]+)$", file.name)
@ -359,12 +356,12 @@ def move():
for file_id in file_ids:
e, file = FileService.get_by_id(file_id)
if not e:
return get_data_error_result(retmsg="File or Folder not found!")
return get_data_error_result(message="File or Folder not found!")
if not file.tenant_id:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
fe, _ = FileService.get_by_id(parent_id)
if not fe:
return get_data_error_result(retmsg="Parent Folder not found!")
return get_data_error_result(message="Parent Folder not found!")
FileService.move_file(file_ids, parent_id)
return get_json_result(data=True)
except Exception as e:

30
api/apps/kb_app.py
View File

@ -47,7 +47,7 @@ def create():
req["created_by"] = current_user.id
e, t = TenantService.get_by_id(current_user.id)
if not e:
return get_data_error_result(retmsg="Tenant not found.")
return get_data_error_result(message="Tenant not found.")
req["embd_id"] = t.embd_id
if not KnowledgebaseService.save(**req):
return get_data_error_result()
@ -65,24 +65,24 @@ def update():
if not KnowledgebaseService.accessible4deletion(req["kb_id"], current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
try:
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)
data=False, message='Only owner of knowledgebase authorized for this operation.', code=RetCode.OPERATING_ERROR)
e, kb = KnowledgebaseService.get_by_id(req["kb_id"])
if not e:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
if req["name"].lower() != kb.name.lower() \
and len(KnowledgebaseService.query(name=req["name"], tenant_id=current_user.id, status=StatusEnum.VALID.value)) > 1:
return get_data_error_result(
retmsg="Duplicated knowledgebase name.")
message="Duplicated knowledgebase name.")
del req["kb_id"]
if not KnowledgebaseService.update_by_id(kb.id, req):
@ -91,7 +91,7 @@ def update():
e, kb = KnowledgebaseService.get_by_id(kb.id)
if not e:
return get_data_error_result(
retmsg="Database error (Knowledgebase rename)!")
message="Database error (Knowledgebase rename)!")
return get_json_result(data=kb.to_json())
except Exception as e:
@ -110,12 +110,12 @@ def detail():
break
else:
return get_json_result(
data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of knowledgebase authorized for this operation.',
code=RetCode.OPERATING_ERROR)
kb = KnowledgebaseService.get_detail(kb_id)
if not kb:
return get_data_error_result(
retmsg="Can't find this knowledgebase!")
message="Can't find this knowledgebase!")
return get_json_result(data=kb)
except Exception as e:
return server_error_response(e)
@ -145,27 +145,27 @@ def rm():
if not KnowledgebaseService.accessible4deletion(req["kb_id"], current_user.id):
return get_json_result(
data=False,
retmsg='No authorization.',
retcode=RetCode.AUTHENTICATION_ERROR
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
try:
kbs = KnowledgebaseService.query(
created_by=current_user.id, id=req["kb_id"])
if not kbs:
return get_json_result(
data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.', retcode=RetCode.OPERATING_ERROR)
data=False, message='Only owner of knowledgebase authorized for this operation.', code=RetCode.OPERATING_ERROR)
for doc in DocumentService.query(kb_id=req["kb_id"]):
if not DocumentService.remove_document(doc, kbs[0].tenant_id):
return get_data_error_result(
retmsg="Database error (Document removal)!")
message="Database error (Document removal)!")
f2d = File2DocumentService.get_by_document_id(doc.id)
FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
File2DocumentService.delete_by_document_id(doc.id)
if not KnowledgebaseService.delete_by_id(req["kb_id"]):
return get_data_error_result(
retmsg="Database error (Knowledgebase removal)!")
message="Database error (Knowledgebase removal)!")
return get_json_result(data=True)
except Exception as e:
return server_error_response(e)

4
api/apps/llm_app.py
View File

@ -98,7 +98,7 @@ def set_api_key():
break
if msg:
return get_data_error_result(retmsg=msg)
return get_data_error_result(message=msg)
llm_config = {
"api_key": req["api_key"],
@ -278,7 +278,7 @@ def add_llm():
pass
if msg:
return get_data_error_result(retmsg=msg)
return get_data_error_result(message=msg)
if not TenantLLMService.filter_update(
[TenantLLM.tenant_id == current_user.id, TenantLLM.llm_factory == factory, TenantLLM.llm_name == llm["llm_name"]], llm):

40
api/apps/sdk/chat.py
View File

@ -32,7 +32,7 @@ def create(tenant_id):
req=request.json
ids= req.get("dataset_ids")
if not ids:
return get_error_data_result(retmsg="`dataset_ids` is required")
return get_error_data_result(message="`dataset_ids` is required")
for kb_id in ids:
kbs = KnowledgebaseService.query(id=kb_id,tenant_id=tenant_id)
if not kbs:
@ -43,7 +43,7 @@ def create(tenant_id):
kbs = KnowledgebaseService.get_by_ids(ids)
embd_count = list(set([kb.embd_id for kb in kbs]))
if len(embd_count) != 1:
return get_result(retmsg='Datasets use different embedding models."',retcode=RetCode.AUTHENTICATION_ERROR)
return get_result(message='Datasets use different embedding models."',code=RetCode.AUTHENTICATION_ERROR)
req["kb_ids"] = ids
# llm
llm = req.get("llm")
@ -55,7 +55,7 @@ def create(tenant_id):
req["llm_setting"] = req.pop("llm")
e, tenant = TenantService.get_by_id(tenant_id)
if not e:
return get_error_data_result(retmsg="Tenant not found!")
return get_error_data_result(message="Tenant not found!")
# prompt
prompt = req.get("prompt")
key_mapping = {"parameters": "variables",
@ -86,12 +86,12 @@ def create(tenant_id):
if not req.get("llm_id"):
req["llm_id"] = tenant.llm_id
if not req.get("name"):
return get_error_data_result(retmsg="`name` is required.")
return get_error_data_result(message="`name` is required.")
if DialogService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg="Duplicated chat name in creating chat.")
return get_error_data_result(message="Duplicated chat name in creating chat.")
# tenant_id
if req.get("tenant_id"):
return get_error_data_result(retmsg="`tenant_id` must not be provided.")
return get_error_data_result(message="`tenant_id` must not be provided.")
req["tenant_id"] = tenant_id
# prompt more parameter
default_prompt = {
@ -117,14 +117,14 @@ def create(tenant_id):
continue
if req['prompt_config']["system"].find("{%s}" % p["key"]) < 0:
return get_error_data_result(
retmsg="Parameter '{}' is not used".format(p["key"]))
message="Parameter '{}' is not used".format(p["key"]))
# save
if not DialogService.save(**req):
return get_error_data_result(retmsg="Fail to new a chat!")
return get_error_data_result(message="Fail to new a chat!")
# response
e, res = DialogService.get_by_id(req["id"])
if not e:
return get_error_data_result(retmsg="Fail to new a chat!")
return get_error_data_result(message="Fail to new a chat!")
res = res.to_json()
renamed_dict = {}
for key, value in res["prompt_config"].items():
@ -150,7 +150,7 @@ def create(tenant_id):
@token_required
def update(tenant_id,chat_id):
if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg='You do not own the chat')
return get_error_data_result(message='You do not own the chat')
req =request.json
ids = req.get("dataset_ids")
if "show_quotation" in req:
@ -170,8 +170,8 @@ def update(tenant_id,chat_id):
embd_count=list(set([kb.embd_id for kb in kbs]))
if len(embd_count) != 1 :
return get_result(
retmsg='Datasets use different embedding models."',
retcode=RetCode.AUTHENTICATION_ERROR)
message='Datasets use different embedding models."',
code=RetCode.AUTHENTICATION_ERROR)
req["kb_ids"] = ids
llm = req.get("llm")
if llm:
@ -182,7 +182,7 @@ def update(tenant_id,chat_id):
req["llm_setting"] = req.pop("llm")
e, tenant = TenantService.get_by_id(tenant_id)
if not e:
return get_error_data_result(retmsg="Tenant not found!")
return get_error_data_result(message="Tenant not found!")
if req.get("rerank_model"):
if not TenantLLMService.query(tenant_id=tenant_id,llm_name=req.get("rerank_model"),model_type="rerank"):
return get_error_data_result(f"`rerank_model` {req.get('rerank_model')} doesn't exist")
@ -207,18 +207,18 @@ def update(tenant_id,chat_id):
res = res.to_json()
if "name" in req:
if not req.get("name"):
return get_error_data_result(retmsg="`name` is not empty.")
return get_error_data_result(message="`name` is not empty.")
if req["name"].lower() != res["name"].lower() \
and len(
DialogService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value)) > 0:
return get_error_data_result(retmsg="Duplicated chat name in updating dataset.")
return get_error_data_result(message="Duplicated chat name in updating dataset.")
if "prompt_config" in req:
res["prompt_config"].update(req["prompt_config"])
for p in res["prompt_config"]["parameters"]:
if p["optional"]:
continue
if res["prompt_config"]["system"].find("{%s}" % p["key"]) < 0:
return get_error_data_result(retmsg="Parameter '{}' is not used".format(p["key"]))
return get_error_data_result(message="Parameter '{}' is not used".format(p["key"]))
if "llm_setting" in req:
res["llm_setting"].update(req["llm_setting"])
req["prompt_config"] = res["prompt_config"]
@ -229,7 +229,7 @@ def update(tenant_id,chat_id):
if "dataset_ids" in req:
req.pop("dataset_ids")
if not DialogService.update_by_id(chat_id, req):
return get_error_data_result(retmsg="Chat not found!")
return get_error_data_result(message="Chat not found!")
return get_result()
@ -250,7 +250,7 @@ def delete(tenant_id):
id_list=ids
for id in id_list:
if not DialogService.query(tenant_id=tenant_id, id=id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg=f"You don't own the chat {id}")
return get_error_data_result(message=f"You don't own the chat {id}")
temp_dict = {"status": StatusEnum.INVALID.value}
DialogService.update_by_id(id, temp_dict)
return get_result()
@ -262,7 +262,7 @@ def list_chat(tenant_id):
name = request.args.get("name")
chat = DialogService.query(id=id,name=name,status=StatusEnum.VALID.value)
if not chat:
return get_error_data_result(retmsg="The chat doesn't exist")
return get_error_data_result(message="The chat doesn't exist")
page_number = int(request.args.get("page", 1))
items_per_page = int(request.args.get("page_size", 1024))
orderby = request.args.get("orderby", "create_time")
@ -302,7 +302,7 @@ def list_chat(tenant_id):
for kb_id in res["kb_ids"]:
kb = KnowledgebaseService.query(id=kb_id)
if not kb :
return get_error_data_result(retmsg=f"Don't exist the kb {kb_id}")
return get_error_data_result(message=f"Don't exist the kb {kb_id}")
kb_list.append(kb[0].to_json())
del res["kb_ids"]
res["datasets"] = kb_list

1070
api/apps/sdk/dataset.py
View File

File diff suppressed because it is too large Load Diff

10
api/apps/sdk/dify_retrieval.py
View File

@ -37,10 +37,10 @@ def retrieval(tenant_id):
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
return build_error_result(error_msg="Knowledgebase not found!", retcode=RetCode.NOT_FOUND)
return build_error_result(message="Knowledgebase not found!", code=RetCode.NOT_FOUND)
if kb.tenant_id != tenant_id:
return build_error_result(error_msg="Knowledgebase not found!", retcode=RetCode.NOT_FOUND)
return build_error_result(message="Knowledgebase not found!", code=RetCode.NOT_FOUND)
embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
@ -71,7 +71,7 @@ def retrieval(tenant_id):
except Exception as e:
if str(e).find("not_found") > 0:
return build_error_result(
error_msg=f'No chunk found! Check the chunk status please!',
retcode=RetCode.NOT_FOUND
message='No chunk found! Check the chunk status please!',
code=RetCode.NOT_FOUND
)
return build_error_result(error_msg=str(e), retcode=RetCode.SERVER_ERROR)
return build_error_result(message=str(e), code=RetCode.SERVER_ERROR)

96
api/apps/sdk/doc.py
View File

@ -110,13 +110,13 @@ def upload(dataset_id, tenant_id):
"""
if "file" not in request.files:
return get_error_data_result(
retmsg="No file part!", retcode=RetCode.ARGUMENT_ERROR
message="No file part!", code=RetCode.ARGUMENT_ERROR
)
file_objs = request.files.getlist("file")
for file_obj in file_objs:
if file_obj.filename == "":
return get_result(
retmsg="No file selected!", retcode=RetCode.ARGUMENT_ERROR
message="No file selected!", code=RetCode.ARGUMENT_ERROR
)
# total size
total_size = 0
@ -127,15 +127,15 @@ def upload(dataset_id, tenant_id):
MAX_TOTAL_FILE_SIZE = 10 * 1024 * 1024
if total_size > MAX_TOTAL_FILE_SIZE:
return get_result(
retmsg=f"Total file size exceeds 10MB limit! ({total_size / (1024 * 1024):.2f} MB)",
retcode=RetCode.ARGUMENT_ERROR,
message=f"Total file size exceeds 10MB limit! ({total_size / (1024 * 1024):.2f} MB)",
code=RetCode.ARGUMENT_ERROR,
)
e, kb = KnowledgebaseService.get_by_id(dataset_id)
if not e:
raise LookupError(f"Can't find the dataset with ID {dataset_id}!")
err, files = FileService.upload_document(kb, file_objs, tenant_id)
if err:
return get_result(retmsg="\n".join(err), retcode=RetCode.SERVER_ERROR)
return get_result(message="\n".join(err), code=RetCode.SERVER_ERROR)
# rename key's name
renamed_doc_list = []
for file in files:
@ -205,20 +205,20 @@ def update_doc(tenant_id, dataset_id, document_id):
"""
req = request.json
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg="You don't own the dataset.")
return get_error_data_result(message="You don't own the dataset.")
doc = DocumentService.query(kb_id=dataset_id, id=document_id)
if not doc:
return get_error_data_result(retmsg="The dataset doesn't own the document.")
return get_error_data_result(message="The dataset doesn't own the document.")
doc = doc[0]
if "chunk_count" in req:
if req["chunk_count"] != doc.chunk_num:
return get_error_data_result(retmsg="Can't change `chunk_count`.")
return get_error_data_result(message="Can't change `chunk_count`.")
if "token_count" in req:
if req["token_count"] != doc.token_num:
return get_error_data_result(retmsg="Can't change `token_count`.")
return get_error_data_result(message="Can't change `token_count`.")
if "progress" in req:
if req["progress"] != doc.progress:
return get_error_data_result(retmsg="Can't change `progress`.")
return get_error_data_result(message="Can't change `progress`.")
if "name" in req and req["name"] != doc.name:
if (
@ -226,16 +226,16 @@ def update_doc(tenant_id, dataset_id, document_id):
!= pathlib.Path(doc.name.lower()).suffix
):
return get_result(
retmsg="The extension of file can't be changed",
retcode=RetCode.ARGUMENT_ERROR,
message="The extension of file can't be changed",
code=RetCode.ARGUMENT_ERROR,
)
for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
if d.name == req["name"]:
return get_error_data_result(
retmsg="Duplicated document name in the same dataset."
message="Duplicated document name in the same dataset."
)
if not DocumentService.update_by_id(document_id, {"name": req["name"]}):
return get_error_data_result(retmsg="Database error (Document rename)!")
return get_error_data_result(message="Database error (Document rename)!")
informs = File2DocumentService.get_by_document_id(document_id)
if informs:
@ -266,7 +266,7 @@ def update_doc(tenant_id, dataset_id, document_id):
return get_result()
if doc.type == FileType.VISUAL or re.search(r"\.(ppt|pptx|pages)$", doc.name):
return get_error_data_result(retmsg="Not supported yet!")
return get_error_data_result(message="Not supported yet!")
e = DocumentService.update_by_id(
doc.id,
@ -278,7 +278,7 @@ def update_doc(tenant_id, dataset_id, document_id):
},
)
if not e:
return get_error_data_result(retmsg="Document not found!")
return get_error_data_result(message="Document not found!")
req["parser_config"] = get_parser_config(
req["chunk_method"], req.get("parser_config")
)
@ -292,7 +292,7 @@ def update_doc(tenant_id, dataset_id, document_id):
doc.process_duation * -1,
)
if not e:
return get_error_data_result(retmsg="Document not found!")
return get_error_data_result(message="Document not found!")
ELASTICSEARCH.deleteByQuery(
Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id)
)
@ -339,11 +339,11 @@ def download(tenant_id, dataset_id, document_id):
type: object
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You do not own the dataset {dataset_id}.")
return get_error_data_result(message=f"You do not own the dataset {dataset_id}.")
doc = DocumentService.query(kb_id=dataset_id, id=document_id)
if not doc:
return get_error_data_result(
retmsg=f"The dataset not own the document {document_id}."
message=f"The dataset not own the document {document_id}."
)
# The process of downloading
doc_id, doc_location = File2DocumentService.get_storage_address(
@ -451,13 +451,13 @@ def list_docs(dataset_id, tenant_id):
description: Processing status.
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}. ")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}. ")
id = request.args.get("id")
name = request.args.get("name")
if not DocumentService.query(id=id, kb_id=dataset_id):
return get_error_data_result(retmsg=f"You don't own the document {id}.")
return get_error_data_result(message=f"You don't own the document {id}.")
if not DocumentService.query(name=name, kb_id=dataset_id):
return get_error_data_result(retmsg=f"You don't own the document {name}.")
return get_error_data_result(message=f"You don't own the document {name}.")
page = int(request.args.get("page", 1))
keywords = request.args.get("keywords", "")
page_size = int(request.args.get("page_size", 1024))
@ -538,7 +538,7 @@ def delete(tenant_id, dataset_id):
type: object
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}. ")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}. ")
req = request.json
if not req:
doc_ids = None
@ -559,16 +559,16 @@ def delete(tenant_id, dataset_id):
try:
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_error_data_result(retmsg="Document not found!")
return get_error_data_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(doc_id)
if not tenant_id:
return get_error_data_result(retmsg="Tenant not found!")
return get_error_data_result(message="Tenant not found!")
b, n = File2DocumentService.get_storage_address(doc_id=doc_id)
if not DocumentService.remove_document(doc, tenant_id):
return get_error_data_result(
retmsg="Database error (Document removal)!"
message="Database error (Document removal)!"
)
f2d = File2DocumentService.get_by_document_id(doc_id)
@ -585,7 +585,7 @@ def delete(tenant_id, dataset_id):
errors += str(e)
if errors:
return get_result(retmsg=errors, retcode=RetCode.SERVER_ERROR)
return get_result(message=errors, code=RetCode.SERVER_ERROR)
return get_result()
@ -630,14 +630,14 @@ def parse(tenant_id, dataset_id):
type: object
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
req = request.json
if not req.get("document_ids"):
return get_error_data_result("`document_ids` is required")
for id in req["document_ids"]:
doc = DocumentService.query(id=id, kb_id=dataset_id)
if not doc:
return get_error_data_result(retmsg=f"You don't own the document {id}.")
return get_error_data_result(message=f"You don't own the document {id}.")
if doc[0].progress != 0.0:
return get_error_data_result(
"Can't stop parsing document with progress at 0 or 100"
@ -699,14 +699,14 @@ def stop_parsing(tenant_id, dataset_id):
type: object
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
req = request.json
if not req.get("document_ids"):
return get_error_data_result("`document_ids` is required")
for id in req["document_ids"]:
doc = DocumentService.query(id=id, kb_id=dataset_id)
if not doc:
return get_error_data_result(retmsg=f"You don't own the document {id}.")
return get_error_data_result(message=f"You don't own the document {id}.")
if int(doc[0].progress) == 1 or int(doc[0].progress) == 0:
return get_error_data_result(
"Can't stop parsing document with progress at 0 or 1"
@ -793,11 +793,11 @@ def list_chunks(tenant_id, dataset_id, document_id):
description: Document details.
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(
retmsg=f"You don't own the document {document_id}."
message=f"You don't own the document {document_id}."
)
doc = doc[0]
req = request.args
@ -965,16 +965,16 @@ def add_chunk(tenant_id, dataset_id, document_id):
description: Important keywords.
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(
retmsg=f"You don't own the document {document_id}."
message=f"You don't own the document {document_id}."
)
doc = doc[0]
req = request.json
if not req.get("content"):
return get_error_data_result(retmsg="`content` is required")
return get_error_data_result(message="`content` is required")
if "important_keywords" in req:
if type(req["important_keywords"]) != list:
return get_error_data_result(
@ -1078,11 +1078,11 @@ def rm_chunk(tenant_id, dataset_id, document_id):
type: object
"""
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(
retmsg=f"You don't own the document {document_id}."
message=f"You don't own the document {document_id}."
)
doc = doc[0]
req = request.json
@ -1104,7 +1104,7 @@ def rm_chunk(tenant_id, dataset_id, document_id):
if not ELASTICSEARCH.deleteByQuery(
Q("ids", values=chunk_list), search.index_name(tenant_id)
):
return get_error_data_result(retmsg="Index updating failure")
return get_error_data_result(message="Index updating failure")
deleted_chunk_ids = chunk_list
chunk_number = len(deleted_chunk_ids)
DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
@ -1170,14 +1170,14 @@ def update_chunk(tenant_id, dataset_id, document_id, chunk_id):
"""
try:
res = ELASTICSEARCH.get(chunk_id, search.index_name(tenant_id))
except Exception as e:
except Exception:
return get_error_data_result(f"Can't find this chunk {chunk_id}")
if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(
retmsg=f"You don't own the document {document_id}."
message=f"You don't own the document {document_id}."
)
doc = doc[0]
query = {
@ -1210,7 +1210,7 @@ def update_chunk(tenant_id, dataset_id, document_id, chunk_id):
arr = [t for t in re.split(r"[\n\t]", d["content_with_weight"]) if len(t) > 1]
if len(arr) != 2:
return get_error_data_result(
retmsg="Q&A must be separated by TAB/ENTER key."
message="Q&A must be separated by TAB/ENTER key."
)
q, a = rmPrefix(arr[0]), rmPrefix(arr[1])
d = beAdoc(
@ -1317,8 +1317,8 @@ def retrieval_test(tenant_id):
embd_nms = list(set([kb.embd_id for kb in kbs]))
if len(embd_nms) != 1:
return get_result(
retmsg='Datasets use different embedding models."',
retcode=RetCode.AUTHENTICATION_ERROR,
message='Datasets use different embedding models."',
code=RetCode.AUTHENTICATION_ERROR,
)
if "question" not in req:
return get_error_data_result("`question` is required.")
@ -1344,7 +1344,7 @@ def retrieval_test(tenant_id):
try:
e, kb = KnowledgebaseService.get_by_id(kb_ids[0])
if not e:
return get_error_data_result(retmsg="Dataset not found!")
return get_error_data_result(message="Dataset not found!")
embd_mdl = TenantLLMService.model_instance(
kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id
)
@ -1398,7 +1398,7 @@ def retrieval_test(tenant_id):
except Exception as e:
if str(e).find("not_found") > 0:
return get_result(
retmsg=f"No chunk found! Check the chunk status please!",
retcode=RetCode.DATA_ERROR,
message="No chunk found! Check the chunk status please!",
code=RetCode.DATA_ERROR,
)
return server_error_response(e)

788
api/apps/sdk/session.py
View File

@ -1,394 +1,394 @@
#
# Copyright 2024 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 json
from functools import partial
from uuid import uuid4
from flask import request, Response
from agent.canvas import Canvas
from api.db import StatusEnum
from api.db.db_models import API4Conversation
from api.db.services.api_service import API4ConversationService
from api.db.services.canvas_service import UserCanvasService
from api.db.services.dialog_service import DialogService, ConversationService, chat
from api.utils import get_uuid
from api.utils.api_utils import get_error_data_result
from api.utils.api_utils import get_result, token_required
@manager.route('/chats/<chat_id>/sessions', methods=['POST'])
@token_required
def create(tenant_id,chat_id):
req = request.json
req["dialog_id"] = chat_id
dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
if not dia:
return get_error_data_result(retmsg="You do not own the assistant.")
conv = {
"id": get_uuid(),
"dialog_id": req["dialog_id"],
"name": req.get("name", "New session"),
"message": [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
}
if not conv.get("name"):
return get_error_data_result(retmsg="`name` can not be empty.")
ConversationService.save(**conv)
e, conv = ConversationService.get_by_id(conv["id"])
if not e:
return get_error_data_result(retmsg="Fail to create a session!")
conv = conv.to_dict()
conv['messages'] = conv.pop("message")
conv["chat_id"] = conv.pop("dialog_id")
del conv["reference"]
return get_result(data=conv)
@manager.route('/agents/<agent_id>/sessions', methods=['POST'])
@token_required
def create_agent_session(tenant_id, agent_id):
req = request.json
e, cvs = UserCanvasService.get_by_id(agent_id)
if not e:
return get_error_data_result("Agent not found.")
if cvs.user_id != tenant_id:
return get_error_data_result(retmsg="You do not own the agent.")
if not isinstance(cvs.dsl, str):
cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
canvas = Canvas(cvs.dsl, tenant_id)
conv = {
"id": get_uuid(),
"dialog_id": cvs.id,
"user_id": req.get("user_id", ""),
"message": [{"role": "assistant", "content": canvas.get_prologue()}],
"source": "agent"
}
API4ConversationService.save(**conv)
return get_result(data=conv)
@manager.route('/chats/<chat_id>/sessions/<session_id>', methods=['PUT'])
@token_required
def update(tenant_id,chat_id,session_id):
req = request.json
req["dialog_id"] = chat_id
conv_id = session_id
conv = ConversationService.query(id=conv_id,dialog_id=chat_id)
if not conv:
return get_error_data_result(retmsg="Session does not exist")
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg="You do not own the session")
if "message" in req or "messages" in req:
return get_error_data_result(retmsg="`message` can not be change")
if "reference" in req:
return get_error_data_result(retmsg="`reference` can not be change")
if "name" in req and not req.get("name"):
return get_error_data_result(retmsg="`name` can not be empty.")
if not ConversationService.update_by_id(conv_id, req):
return get_error_data_result(retmsg="Session updates error")
return get_result()
@manager.route('/chats/<chat_id>/completions', methods=['POST'])
@token_required
def completion(tenant_id, chat_id):
req = request.json
if not req.get("session_id"):
conv = {
"id": get_uuid(),
"dialog_id": chat_id,
"name": req.get("name", "New session"),
"message": [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
}
if not conv.get("name"):
return get_error_data_result(retmsg="`name` can not be empty.")
ConversationService.save(**conv)
e, conv = ConversationService.get_by_id(conv["id"])
session_id=conv.id
else:
session_id = req.get("session_id")
if not req.get("question"):
return get_error_data_result(retmsg="Please input your question.")
conv = ConversationService.query(id=session_id,dialog_id=chat_id)
if not conv:
return get_error_data_result(retmsg="Session does not exist")
conv = conv[0]
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg="You do not own the chat")
msg = []
question = {
"content": req.get("question"),
"role": "user",
"id": str(uuid4())
}
conv.message.append(question)
for m in conv.message:
if m["role"] == "system": continue
if m["role"] == "assistant" and not msg: continue
msg.append(m)
message_id = msg[-1].get("id")
e, dia = DialogService.get_by_id(conv.dialog_id)
if not conv.reference:
conv.reference = []
conv.message.append({"role": "assistant", "content": "", "id": message_id})
conv.reference.append({"chunks": [], "doc_aggs": []})
def fillin_conv(ans):
nonlocal conv, message_id
if not conv.reference:
conv.reference.append(ans["reference"])
else:
conv.reference[-1] = ans["reference"]
conv.message[-1] = {"role": "assistant", "content": ans["answer"],
"id": message_id, "prompt": ans.get("prompt", "")}
ans["id"] = message_id
ans["session_id"]=session_id
def stream():
nonlocal dia, msg, req, conv
try:
for ans in chat(dia, msg, **req):
fillin_conv(ans)
yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
ConversationService.update_by_id(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e),"reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
if req.get("stream", True):
resp = Response(stream(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
resp.headers.add_header("Connection", "keep-alive")
resp.headers.add_header("X-Accel-Buffering", "no")
resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
return resp
else:
answer = None
for ans in chat(dia, msg, **req):
answer = ans
fillin_conv(ans)
ConversationService.update_by_id(conv.id, conv.to_dict())
break
return get_result(data=answer)
@manager.route('/agents/<agent_id>/completions', methods=['POST'])
@token_required
def agent_completion(tenant_id, agent_id):
req = request.json
e, cvs = UserCanvasService.get_by_id(agent_id)
if not e:
return get_error_data_result("Agent not found.")
if cvs.user_id != tenant_id:
return get_error_data_result(retmsg="You do not own the agent.")
if not isinstance(cvs.dsl, str):
cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
canvas = Canvas(cvs.dsl, tenant_id)
msg = []
for m in req["messages"]:
if m["role"] == "system":
continue
if m["role"] == "assistant" and not msg:
continue
msg.append(m)
if not msg[-1].get("id"): msg[-1]["id"] = get_uuid()
message_id = msg[-1]["id"]
if not req.get("session_id"):
session_id = get_uuid()
conv = {
"id": session_id,
"dialog_id": cvs.id,
"user_id": req.get("user_id", ""),
"message": [{"role": "assistant", "content": canvas.get_prologue()}],
"source": "agent"
}
API4ConversationService.save(**conv)
conv = API4Conversation(**conv)
else:
session_id = req.get("session_id")
e, conv = API4ConversationService.get_by_id(req["session_id"])
if not e:
return get_error_data_result(retmsg="Session not found!")
if "quote" not in req: req["quote"] = False
stream = req.get("stream", True)
def fillin_conv(ans):
nonlocal conv, message_id
if not conv.reference:
conv.reference.append(ans["reference"])
else:
conv.reference[-1] = ans["reference"]
conv.message[-1] = {"role": "assistant", "content": ans["answer"], "id": message_id}
ans["id"] = message_id
ans["session_id"] = session_id
def rename_field(ans):
reference = ans['reference']
if not isinstance(reference, dict):
return
for chunk_i in reference.get('chunks', []):
if 'docnm_kwd' in chunk_i:
chunk_i['doc_name'] = chunk_i['docnm_kwd']
chunk_i.pop('docnm_kwd')
conv.message.append(msg[-1])
if not conv.reference:
conv.reference = []
conv.message.append({"role": "assistant", "content": "", "id": message_id})
conv.reference.append({"chunks": [], "doc_aggs": []})
final_ans = {"reference": [], "content": ""}
canvas.messages.append(msg[-1])
canvas.add_user_input(msg[-1]["content"])
answer = canvas.run(stream=stream)
assert answer is not None, "Nothing. Is it over?"
if stream:
assert isinstance(answer, partial), "Nothing. Is it over?"
def sse():
nonlocal answer, cvs, conv
try:
for ans in answer():
for k in ans.keys():
final_ans[k] = ans[k]
ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
fillin_conv(ans)
rename_field(ans)
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
ensure_ascii=False) + "\n\n"
canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
if final_ans.get("reference"):
canvas.reference.append(final_ans["reference"])
cvs.dsl = json.loads(str(canvas))
API4ConversationService.append_message(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(sse(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
resp.headers.add_header("Connection", "keep-alive")
resp.headers.add_header("X-Accel-Buffering", "no")
resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
return resp
final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
if final_ans.get("reference"):
canvas.reference.append(final_ans["reference"])
cvs.dsl = json.loads(str(canvas))
result = {"answer": final_ans["content"], "reference": final_ans.get("reference", [])}
fillin_conv(result)
API4ConversationService.append_message(conv.id, conv.to_dict())
rename_field(result)
return get_result(data=result)
@manager.route('/chats/<chat_id>/sessions', methods=['GET'])
@token_required
def list(chat_id,tenant_id):
if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg=f"You don't own the assistant {chat_id}.")
id = request.args.get("id")
name = request.args.get("name")
page_number = int(request.args.get("page", 1))
items_per_page = int(request.args.get("page_size", 1024))
orderby = request.args.get("orderby", "create_time")
if request.args.get("desc") == "False" or request.args.get("desc") == "false":
desc = False
else:
desc = True
convs = ConversationService.get_list(chat_id,page_number,items_per_page,orderby,desc,id,name)
if not convs:
return get_result(data=[])
for conv in convs:
conv['messages'] = conv.pop("message")
infos = conv["messages"]
for info in infos:
if "prompt" in info:
info.pop("prompt")
conv["chat"] = conv.pop("dialog_id")
if conv["reference"]:
messages = conv["messages"]
message_num = 0
chunk_num = 0
while message_num < len(messages):
if message_num != 0 and messages[message_num]["role"] != "user":
chunk_list = []
if "chunks" in conv["reference"][chunk_num]:
chunks = conv["reference"][chunk_num]["chunks"]
for chunk in chunks:
new_chunk = {
"id": chunk["chunk_id"],
"content": chunk["content_with_weight"],
"document_id": chunk["doc_id"],
"document_name": chunk["docnm_kwd"],
"dataset_id": chunk["kb_id"],
"image_id": chunk["img_id"],
"similarity": chunk["similarity"],
"vector_similarity": chunk["vector_similarity"],
"term_similarity": chunk["term_similarity"],
"positions": chunk["positions"],
}
chunk_list.append(new_chunk)
chunk_num += 1
messages[message_num]["reference"] = chunk_list
message_num += 1
del conv["reference"]
return get_result(data=convs)
@manager.route('/chats/<chat_id>/sessions', methods=["DELETE"])
@token_required
def delete(tenant_id,chat_id):
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg="You don't own the chat")
req = request.json
convs = ConversationService.query(dialog_id=chat_id)
if not req:
ids = None
else:
ids=req.get("ids")
if not ids:
conv_list = []
for conv in convs:
conv_list.append(conv.id)
else:
conv_list=ids
for id in conv_list:
conv = ConversationService.query(id=id,dialog_id=chat_id)
if not conv:
return get_error_data_result(retmsg="The chat doesn't own the session")
ConversationService.delete_by_id(id)
return get_result()
#
# Copyright 2024 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 json
from functools import partial
from uuid import uuid4
from flask import request, Response
from agent.canvas import Canvas
from api.db import StatusEnum
from api.db.db_models import API4Conversation
from api.db.services.api_service import API4ConversationService
from api.db.services.canvas_service import UserCanvasService
from api.db.services.dialog_service import DialogService, ConversationService, chat
from api.utils import get_uuid
from api.utils.api_utils import get_error_data_result
from api.utils.api_utils import get_result, token_required
@manager.route('/chats/<chat_id>/sessions', methods=['POST'])
@token_required
def create(tenant_id,chat_id):
req = request.json
req["dialog_id"] = chat_id
dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
if not dia:
return get_error_data_result(message="You do not own the assistant.")
conv = {
"id": get_uuid(),
"dialog_id": req["dialog_id"],
"name": req.get("name", "New session"),
"message": [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
}
if not conv.get("name"):
return get_error_data_result(message="`name` can not be empty.")
ConversationService.save(**conv)
e, conv = ConversationService.get_by_id(conv["id"])
if not e:
return get_error_data_result(message="Fail to create a session!")
conv = conv.to_dict()
conv['messages'] = conv.pop("message")
conv["chat_id"] = conv.pop("dialog_id")
del conv["reference"]
return get_result(data=conv)
@manager.route('/agents/<agent_id>/sessions', methods=['POST'])
@token_required
def create_agent_session(tenant_id, agent_id):
req = request.json
e, cvs = UserCanvasService.get_by_id(agent_id)
if not e:
return get_error_data_result("Agent not found.")
if cvs.user_id != tenant_id:
return get_error_data_result(message="You do not own the agent.")
if not isinstance(cvs.dsl, str):
cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
canvas = Canvas(cvs.dsl, tenant_id)
conv = {
"id": get_uuid(),
"dialog_id": cvs.id,
"user_id": req.get("user_id", ""),
"message": [{"role": "assistant", "content": canvas.get_prologue()}],
"source": "agent"
}
API4ConversationService.save(**conv)
return get_result(data=conv)
@manager.route('/chats/<chat_id>/sessions/<session_id>', methods=['PUT'])
@token_required
def update(tenant_id,chat_id,session_id):
req = request.json
req["dialog_id"] = chat_id
conv_id = session_id
conv = ConversationService.query(id=conv_id,dialog_id=chat_id)
if not conv:
return get_error_data_result(message="Session does not exist")
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(message="You do not own the session")
if "message" in req or "messages" in req:
return get_error_data_result(message="`message` can not be change")
if "reference" in req:
return get_error_data_result(message="`reference` can not be change")
if "name" in req and not req.get("name"):
return get_error_data_result(message="`name` can not be empty.")
if not ConversationService.update_by_id(conv_id, req):
return get_error_data_result(message="Session updates error")
return get_result()
@manager.route('/chats/<chat_id>/completions', methods=['POST'])
@token_required
def completion(tenant_id, chat_id):
req = request.json
if not req.get("session_id"):
conv = {
"id": get_uuid(),
"dialog_id": chat_id,
"name": req.get("name", "New session"),
"message": [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
}
if not conv.get("name"):
return get_error_data_result(message="`name` can not be empty.")
ConversationService.save(**conv)
e, conv = ConversationService.get_by_id(conv["id"])
session_id=conv.id
else:
session_id = req.get("session_id")
if not req.get("question"):
return get_error_data_result(message="Please input your question.")
conv = ConversationService.query(id=session_id,dialog_id=chat_id)
if not conv:
return get_error_data_result(message="Session does not exist")
conv = conv[0]
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(message="You do not own the chat")
msg = []
question = {
"content": req.get("question"),
"role": "user",
"id": str(uuid4())
}
conv.message.append(question)
for m in conv.message:
if m["role"] == "system": continue
if m["role"] == "assistant" and not msg: continue
msg.append(m)
message_id = msg[-1].get("id")
e, dia = DialogService.get_by_id(conv.dialog_id)
if not conv.reference:
conv.reference = []
conv.message.append({"role": "assistant", "content": "", "id": message_id})
conv.reference.append({"chunks": [], "doc_aggs": []})
def fillin_conv(ans):
nonlocal conv, message_id
if not conv.reference:
conv.reference.append(ans["reference"])
else:
conv.reference[-1] = ans["reference"]
conv.message[-1] = {"role": "assistant", "content": ans["answer"],
"id": message_id, "prompt": ans.get("prompt", "")}
ans["id"] = message_id
ans["session_id"]=session_id
def stream():
nonlocal dia, msg, req, conv
try:
for ans in chat(dia, msg, **req):
fillin_conv(ans)
yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
ConversationService.update_by_id(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e),"reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
if req.get("stream", True):
resp = Response(stream(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
resp.headers.add_header("Connection", "keep-alive")
resp.headers.add_header("X-Accel-Buffering", "no")
resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
return resp
else:
answer = None
for ans in chat(dia, msg, **req):
answer = ans
fillin_conv(ans)
ConversationService.update_by_id(conv.id, conv.to_dict())
break
return get_result(data=answer)
@manager.route('/agents/<agent_id>/completions', methods=['POST'])
@token_required
def agent_completion(tenant_id, agent_id):
req = request.json
e, cvs = UserCanvasService.get_by_id(agent_id)
if not e:
return get_error_data_result("Agent not found.")
if cvs.user_id != tenant_id:
return get_error_data_result(message="You do not own the agent.")
if not isinstance(cvs.dsl, str):
cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
canvas = Canvas(cvs.dsl, tenant_id)
msg = []
for m in req["messages"]:
if m["role"] == "system":
continue
if m["role"] == "assistant" and not msg:
continue
msg.append(m)
if not msg[-1].get("id"): msg[-1]["id"] = get_uuid()
message_id = msg[-1]["id"]
if not req.get("session_id"):
session_id = get_uuid()
conv = {
"id": session_id,
"dialog_id": cvs.id,
"user_id": req.get("user_id", ""),
"message": [{"role": "assistant", "content": canvas.get_prologue()}],
"source": "agent"
}
API4ConversationService.save(**conv)
conv = API4Conversation(**conv)
else:
session_id = req.get("session_id")
e, conv = API4ConversationService.get_by_id(req["session_id"])
if not e:
return get_error_data_result(message="Session not found!")
if "quote" not in req: req["quote"] = False
stream = req.get("stream", True)
def fillin_conv(ans):
nonlocal conv, message_id
if not conv.reference:
conv.reference.append(ans["reference"])
else:
conv.reference[-1] = ans["reference"]
conv.message[-1] = {"role": "assistant", "content": ans["answer"], "id": message_id}
ans["id"] = message_id
ans["session_id"] = session_id
def rename_field(ans):
reference = ans['reference']
if not isinstance(reference, dict):
return
for chunk_i in reference.get('chunks', []):
if 'docnm_kwd' in chunk_i:
chunk_i['doc_name'] = chunk_i['docnm_kwd']
chunk_i.pop('docnm_kwd')
conv.message.append(msg[-1])
if not conv.reference:
conv.reference = []
conv.message.append({"role": "assistant", "content": "", "id": message_id})
conv.reference.append({"chunks": [], "doc_aggs": []})
final_ans = {"reference": [], "content": ""}
canvas.messages.append(msg[-1])
canvas.add_user_input(msg[-1]["content"])
answer = canvas.run(stream=stream)
assert answer is not None, "Nothing. Is it over?"
if stream:
assert isinstance(answer, partial), "Nothing. Is it over?"
def sse():
nonlocal answer, cvs, conv
try:
for ans in answer():
for k in ans.keys():
final_ans[k] = ans[k]
ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
fillin_conv(ans)
rename_field(ans)
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans},
ensure_ascii=False) + "\n\n"
canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
if final_ans.get("reference"):
canvas.reference.append(final_ans["reference"])
cvs.dsl = json.loads(str(canvas))
API4ConversationService.append_message(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"code": 500, "message": str(e),
"data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(sse(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
resp.headers.add_header("Connection", "keep-alive")
resp.headers.add_header("X-Accel-Buffering", "no")
resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
return resp
final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
if final_ans.get("reference"):
canvas.reference.append(final_ans["reference"])
cvs.dsl = json.loads(str(canvas))
result = {"answer": final_ans["content"], "reference": final_ans.get("reference", [])}
fillin_conv(result)
API4ConversationService.append_message(conv.id, conv.to_dict())
rename_field(result)
return get_result(data=result)
@manager.route('/chats/<chat_id>/sessions', methods=['GET'])
@token_required
def list(chat_id,tenant_id):
if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
return get_error_data_result(message=f"You don't own the assistant {chat_id}.")
id = request.args.get("id")
name = request.args.get("name")
page_number = int(request.args.get("page", 1))
items_per_page = int(request.args.get("page_size", 1024))
orderby = request.args.get("orderby", "create_time")
if request.args.get("desc") == "False" or request.args.get("desc") == "false":
desc = False
else:
desc = True
convs = ConversationService.get_list(chat_id,page_number,items_per_page,orderby,desc,id,name)
if not convs:
return get_result(data=[])
for conv in convs:
conv['messages'] = conv.pop("message")
infos = conv["messages"]
for info in infos:
if "prompt" in info:
info.pop("prompt")
conv["chat"] = conv.pop("dialog_id")
if conv["reference"]:
messages = conv["messages"]
message_num = 0
chunk_num = 0
while message_num < len(messages):
if message_num != 0 and messages[message_num]["role"] != "user":
chunk_list = []
if "chunks" in conv["reference"][chunk_num]:
chunks = conv["reference"][chunk_num]["chunks"]
for chunk in chunks:
new_chunk = {
"id": chunk["chunk_id"],
"content": chunk["content_with_weight"],
"document_id": chunk["doc_id"],
"document_name": chunk["docnm_kwd"],
"dataset_id": chunk["kb_id"],
"image_id": chunk["img_id"],
"similarity": chunk["similarity"],
"vector_similarity": chunk["vector_similarity"],
"term_similarity": chunk["term_similarity"],
"positions": chunk["positions"],
}
chunk_list.append(new_chunk)
chunk_num += 1
messages[message_num]["reference"] = chunk_list
message_num += 1
del conv["reference"]
return get_result(data=convs)
@manager.route('/chats/<chat_id>/sessions', methods=["DELETE"])
@token_required
def delete(tenant_id,chat_id):
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(message="You don't own the chat")
req = request.json
convs = ConversationService.query(dialog_id=chat_id)
if not req:
ids = None
else:
ids=req.get("ids")
if not ids:
conv_list = []
for conv in convs:
conv_list.append(conv.id)
else:
conv_list=ids
for id in conv_list:
conv = ConversationService.query(id=id,dialog_id=chat_id)
if not conv:
return get_error_data_result(message="The chat doesn't own the session")
ConversationService.delete_by_id(id)
return get_result()

8
api/apps/system_app.py
View File

@ -29,8 +29,6 @@ from api.utils.api_utils import (
get_data_error_result,
server_error_response,
generate_confirmation_token,
request,
validate_request,
)
from api.versions import get_rag_version
from rag.utils.es_conn import ELASTICSEARCH
@ -209,7 +207,7 @@ def new_token():
try:
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
tenant_id = tenants[0].tenant_id
obj = {
@ -222,7 +220,7 @@ def new_token():
}
if not APITokenService.save(**obj):
return get_data_error_result(retmsg="Fail to new a dialog!")
return get_data_error_result(message="Fail to new a dialog!")
return get_json_result(data=obj)
except Exception as e:
@ -263,7 +261,7 @@ def token_list():
try:
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
objs = APITokenService.query(tenant_id=tenants[0].tenant_id)
return get_json_result(data=[o.to_dict() for o in objs])

6
api/apps/tenant_app.py
View File

@ -44,14 +44,14 @@ def create(tenant_id):
req = request.json
usrs = UserService.query(email=req["email"])
if not usrs:
return get_data_error_result(retmsg="User not found.")
return get_data_error_result(message="User not found.")
user_id = usrs[0].id
user_tenants = UserTenantService.query(user_id=user_id, tenant_id=tenant_id)
if user_tenants:
if user_tenants[0].status == UserTenantRole.NORMAL.value:
return get_data_error_result(retmsg="This user is in the team already.")
return get_data_error_result(retmsg="Invitation notification is sent.")
return get_data_error_result(message="This user is in the team already.")
return get_data_error_result(message="Invitation notification is sent.")
UserTenantService.save(
id=get_uuid(),

48
api/apps/user_app.py
View File

@ -36,7 +36,7 @@ from api.utils import (
current_timestamp,
datetime_format,
)
from api.db import UserTenantRole, LLMType, FileType
from api.db import UserTenantRole, FileType
from api.settings import (
RetCode,
GITHUB_OAUTH,
@ -90,7 +90,7 @@ def login():
"""
if not request.json:
return get_json_result(
data=False, retcode=RetCode.AUTHENTICATION_ERROR, retmsg="Unauthorized!"
data=False, code=RetCode.AUTHENTICATION_ERROR, message="Unauthorized!"
)
email = request.json.get("email", "")
@ -98,8 +98,8 @@ def login():
if not users:
return get_json_result(
data=False,
retcode=RetCode.AUTHENTICATION_ERROR,
retmsg=f"Email: {email} is not registered!",
code=RetCode.AUTHENTICATION_ERROR,
message=f"Email: {email} is not registered!",
)
password = request.json.get("password")
@ -107,7 +107,7 @@ def login():
password = decrypt(password)
except BaseException:
return get_json_result(
data=False, retcode=RetCode.SERVER_ERROR, retmsg="Fail to crypt password"
data=False, code=RetCode.SERVER_ERROR, message="Fail to crypt password"
)
user = UserService.query_user(email, password)
@ -119,12 +119,12 @@ def login():
user.update_date = (datetime_format(datetime.now()),)
user.save()
msg = "Welcome back!"
return construct_response(data=response_data, auth=user.get_id(), retmsg=msg)
return construct_response(data=response_data, auth=user.get_id(), message=msg)
else:
return get_json_result(
data=False,
retcode=RetCode.AUTHENTICATION_ERROR,
retmsg="Email and password do not match!",
code=RetCode.AUTHENTICATION_ERROR,
message="Email and password do not match!",
)
@ -323,7 +323,7 @@ def user_info_from_feishu(access_token):
"Authorization": f"Bearer {access_token}",
}
res = requests.get(
f"https://open.feishu.cn/open-apis/authen/v1/user_info", headers=headers
"https://open.feishu.cn/open-apis/authen/v1/user_info", headers=headers
)
user_info = res.json()["data"]
user_info["email"] = None if user_info.get("email") == "" else user_info["email"]
@ -409,8 +409,8 @@ def setting_user():
):
return get_json_result(
data=False,
retcode=RetCode.AUTHENTICATION_ERROR,
retmsg="Password error!",
code=RetCode.AUTHENTICATION_ERROR,
message="Password error!",
)
if new_password:
@ -438,7 +438,7 @@ def setting_user():
except Exception as e:
stat_logger.exception(e)
return get_json_result(
data=False, retmsg="Update failure!", retcode=RetCode.EXCEPTION_ERROR
data=False, message="Update failure!", code=RetCode.EXCEPTION_ERROR
)
@ -474,21 +474,21 @@ def user_profile():
def rollback_user_registration(user_id):
try:
UserService.delete_by_id(user_id)
except Exception as e:
except Exception:
pass
try:
TenantService.delete_by_id(user_id)
except Exception as e:
except Exception:
pass
try:
u = UserTenantService.query(tenant_id=user_id)
if u:
UserTenantService.delete_by_id(u[0].id)
except Exception as e:
except Exception:
pass
try:
TenantLLM.delete().where(TenantLLM.tenant_id == user_id).execute()
except Exception as e:
except Exception:
pass
@ -581,16 +581,16 @@ def user_add():
if not re.match(r"^[\w\._-]+@([\w_-]+\.)+[\w-]{2,5}$", email_address):
return get_json_result(
data=False,
retmsg=f"Invalid email address: {email_address}!",
retcode=RetCode.OPERATING_ERROR,
message=f"Invalid email address: {email_address}!",
code=RetCode.OPERATING_ERROR,
)
# Check if the email address is already used
if UserService.query(email=email_address):
return get_json_result(
data=False,
retmsg=f"Email: {email_address} has already registered!",
retcode=RetCode.OPERATING_ERROR,
message=f"Email: {email_address} has already registered!",
code=RetCode.OPERATING_ERROR,
)
# Construct user info data
@ -617,15 +617,15 @@ def user_add():
return construct_response(
data=user.to_json(),
auth=user.get_id(),
retmsg=f"{nickname}, welcome aboard!",
message=f"{nickname}, welcome aboard!",
)
except Exception as e:
rollback_user_registration(user_id)
stat_logger.exception(e)
return get_json_result(
data=False,
retmsg=f"User registration failure, error: {str(e)}",
retcode=RetCode.EXCEPTION_ERROR,
message=f"User registration failure, error: {str(e)}",
code=RetCode.EXCEPTION_ERROR,
)
@ -661,7 +661,7 @@ def tenant_info():
try:
tenants = TenantService.get_info_by(current_user.id)
if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_data_error_result(message="Tenant not found!")
return get_json_result(data=tenants[0])
except Exception as e:
return server_error_response(e)