add API for tenant function (#2866)

### What problem does this PR solve?

feat: API access key management
https://github.com/infiniflow/ragflow/issues/2846
feat: Render markdown file with remark-loader
https://github.com/infiniflow/ragflow/issues/2846

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Kevin Hu 2024-10-16 16:10:24 +08:00 committed by GitHub
parent e5d3ab0332
commit b9fa00f341
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 96 additions and 64 deletions

View File

@ -18,8 +18,6 @@ from functools import partial
from flask import request, Response from flask import request, Response
from flask_login import login_required, current_user from flask_login import login_required, current_user
from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
from api.db.services.dialog_service import full_question
from api.db.services.user_service import TenantService
from api.settings import RetCode from api.settings import RetCode
from api.utils import get_uuid from api.utils import get_uuid
from api.utils.api_utils import get_json_result, server_error_response, validate_request, get_data_error_result from api.utils.api_utils import get_json_result, server_error_response, validate_request, get_data_error_result
@ -111,8 +109,9 @@ def run():
if "message" in req: if "message" in req:
canvas.messages.append({"role": "user", "content": req["message"], "id": message_id}) canvas.messages.append({"role": "user", "content": req["message"], "id": message_id})
if len([m for m in canvas.messages if m["role"] == "user"]) > 1: if len([m for m in canvas.messages if m["role"] == "user"]) > 1:
ten = TenantService.get_by_user_id(current_user.id)[0] #ten = TenantService.get_info_by(current_user.id)[0]
#req["message"] = full_question(ten["tenant_id"], ten["llm_id"], canvas.messages) #req["message"] = full_question(ten["tenant_id"], ten["llm_id"], canvas.messages)
pass
canvas.add_user_input(req["message"]) canvas.add_user_input(req["message"])
answer = canvas.run(stream=stream) answer = canvas.run(stream=stream)
print(canvas) print(canvas)

View File

@ -218,7 +218,7 @@ def tts():
req = request.json req = request.json
text = req["text"] text = req["text"]
tenants = TenantService.get_by_user_id(current_user.id) tenants = TenantService.get_info_by(current_user.id)
if not tenants: if not tenants:
return get_data_error_result(retmsg="Tenant not found!") return get_data_error_result(retmsg="Tenant not found!")

View File

@ -133,15 +133,9 @@ def token_list():
return server_error_response(e) return server_error_response(e)
@manager.route('/rm', methods=['POST']) @manager.route('/token/<token>', methods=['DELETE'])
@validate_request("tokens", "tenant_id")
@login_required @login_required
def rm(): def rm(token):
req = request.json
try:
for token in req["tokens"]:
APITokenService.filter_delete( APITokenService.filter_delete(
[APIToken.tenant_id == req["tenant_id"], APIToken.token == token]) [APIToken.tenant_id == current_user.id, APIToken.token == token])
return get_json_result(data=True) return get_json_result(data=True)
except Exception as e:
return server_error_response(e)

View File

@ -15,25 +15,14 @@
# #
from flask import request from flask import request
from flask_login import current_user, login_required from flask_login import login_required, current_user
from api.db import UserTenantRole, StatusEnum from api.db import UserTenantRole, StatusEnum
from api.db.db_models import UserTenant from api.db.db_models import UserTenant
from api.db.services.user_service import TenantService, UserTenantService from api.db.services.user_service import UserTenantService, UserService
from api.settings import RetCode
from api.utils import get_uuid from api.utils import get_uuid, delta_seconds
from api.utils.api_utils import get_json_result, validate_request, server_error_response from api.utils.api_utils import get_json_result, validate_request, server_error_response, get_data_error_result
@manager.route("/list", methods=["GET"])
@login_required
def tenant_list():
try:
tenants = TenantService.get_by_user_id(current_user.id)
return get_json_result(data=tenants)
except Exception as e:
return server_error_response(e)
@manager.route("/<tenant_id>/user/list", methods=["GET"]) @manager.route("/<tenant_id>/user/list", methods=["GET"])
@ -41,6 +30,8 @@ def tenant_list():
def user_list(tenant_id): def user_list(tenant_id):
try: try:
users = UserTenantService.get_by_tenant_id(tenant_id) users = UserTenantService.get_by_tenant_id(tenant_id)
for u in users:
u["delta_seconds"] = delta_seconds(u["update_date"])
return get_json_result(data=users) return get_json_result(data=users)
except Exception as e: except Exception as e:
return server_error_response(e) return server_error_response(e)
@ -48,30 +39,31 @@ def user_list(tenant_id):
@manager.route('/<tenant_id>/user', methods=['POST']) @manager.route('/<tenant_id>/user', methods=['POST'])
@login_required @login_required
@validate_request("user_id") @validate_request("email")
def create(tenant_id): def create(tenant_id):
user_id = request.json.get("user_id") req = request.json
if not user_id: usrs = UserService.query(email=req["email"])
return get_json_result( if not usrs:
data=False, retmsg='Lack of "USER ID"', retcode=RetCode.ARGUMENT_ERROR) return get_data_error_result(retmsg="User not found.")
try: user_id = usrs[0].id
user_tenants = UserTenantService.query(user_id=user_id, tenant_id=tenant_id) user_tenants = UserTenantService.query(user_id=user_id, tenant_id=tenant_id)
if user_tenants: if user_tenants:
uuid = user_tenants[0].id if user_tenants[0].status == UserTenantRole.NORMAL.value:
return get_json_result(data={"id": uuid}) return get_data_error_result(retmsg="This user is in the team already.")
return get_data_error_result(retmsg="Invitation notification is sent.")
uuid = get_uuid()
UserTenantService.save( UserTenantService.save(
id = uuid, id=get_uuid(),
user_id = user_id, user_id=user_id,
tenant_id = tenant_id, tenant_id=tenant_id,
role = UserTenantRole.NORMAL.value, role=UserTenantRole.INVITE,
status = StatusEnum.VALID.value) status=StatusEnum.VALID.value)
return get_json_result(data={"id": uuid}) usr = list(usrs.dicts())[0]
except Exception as e: usr = {k: v for k, v in usr.items() if k in ["id", "avatar", "email", "nickname"]}
return server_error_response(e)
return get_json_result(data=usr)
@manager.route('/<tenant_id>/user/<user_id>', methods=['DELETE']) @manager.route('/<tenant_id>/user/<user_id>', methods=['DELETE'])
@ -83,3 +75,24 @@ def rm(tenant_id, user_id):
except Exception as e: except Exception as e:
return server_error_response(e) return server_error_response(e)
@manager.route("/list", methods=["GET"])
@login_required
def tenant_list():
try:
users = UserTenantService.get_tenants_by_user_id(current_user.id)
for u in users:
u["delta_seconds"] = delta_seconds(u["update_date"])
return get_json_result(data=users)
except Exception as e:
return server_error_response(e)
@manager.route("/agree/<tenant_id>", methods=["GET"])
@login_required
def agree(tenant_id):
try:
UserTenantService.filter_update([UserTenant.tenant_id == tenant_id, UserTenant.user_id == current_user.id], {"role": UserTenantRole.NORMAL})
return get_json_result(data=True)
except Exception as e:
return server_error_response(e)

View File

@ -23,7 +23,7 @@ from flask_login import login_required, current_user, login_user, logout_user
from api.db.db_models import TenantLLM from api.db.db_models import TenantLLM
from api.db.services.llm_service import TenantLLMService, LLMService from api.db.services.llm_service import TenantLLMService, LLMService
from api.utils.api_utils import server_error_response, validate_request from api.utils.api_utils import server_error_response, validate_request, get_data_error_result
from api.utils import get_uuid, get_format_time, decrypt, download_img, current_timestamp, datetime_format from api.utils import get_uuid, get_format_time, decrypt, download_img, current_timestamp, datetime_format
from api.db import UserTenantRole, LLMType, FileType from api.db import UserTenantRole, LLMType, FileType
from api.settings import RetCode, GITHUB_OAUTH, FEISHU_OAUTH, CHAT_MDL, EMBEDDING_MDL, ASR_MDL, IMAGE2TEXT_MDL, PARSERS, \ from api.settings import RetCode, GITHUB_OAUTH, FEISHU_OAUTH, CHAT_MDL, EMBEDDING_MDL, ASR_MDL, IMAGE2TEXT_MDL, PARSERS, \
@ -402,8 +402,10 @@ def user_add():
@login_required @login_required
def tenant_info(): def tenant_info():
try: try:
tenants = TenantService.get_by_user_id(current_user.id)[0] tenants = TenantService.get_info_by(current_user.id)
return get_json_result(data=tenants) if not tenants:
return get_data_error_result(retmsg="Tenant not found!")
return get_json_result(data=tenants[0])
except Exception as e: except Exception as e:
return server_error_response(e) return server_error_response(e)

View File

@ -27,6 +27,7 @@ class UserTenantRole(StrEnum):
OWNER = 'owner' OWNER = 'owner'
ADMIN = 'admin' ADMIN = 'admin'
NORMAL = 'normal' NORMAL = 'normal'
INVITE = 'invite'
class TenantPermission(StrEnum): class TenantPermission(StrEnum):

View File

@ -362,7 +362,7 @@ class DocumentService(CommonService):
elif finished: elif finished:
if d["parser_config"].get("raptor", {}).get("use_raptor") and d["progress_msg"].lower().find(" raptor")<0: if d["parser_config"].get("raptor", {}).get("use_raptor") and d["progress_msg"].lower().find(" raptor")<0:
queue_raptor_tasks(d) queue_raptor_tasks(d)
prg *= 0.98 prg = 0.98 * len(tsks)/(len(tsks)+1)
msg.append("------ RAPTOR -------") msg.append("------ RAPTOR -------")
else: else:
status = TaskStatus.DONE.value status = TaskStatus.DONE.value
@ -379,6 +379,7 @@ class DocumentService(CommonService):
info["progress_msg"] = msg info["progress_msg"] = msg
cls.update_by_id(d["id"], info) cls.update_by_id(d["id"], info)
except Exception as e: except Exception as e:
if str(e).find("'0'") < 0:
stat_logger.error("fetch task exception:" + str(e)) stat_logger.error("fetch task exception:" + str(e))
@classmethod @classmethod

View File

@ -87,7 +87,7 @@ class TenantService(CommonService):
@classmethod @classmethod
@DB.connection_context() @DB.connection_context()
def get_by_user_id(cls, user_id): def get_info_by(cls, user_id):
fields = [ fields = [
cls.model.id.alias("tenant_id"), cls.model.id.alias("tenant_id"),
cls.model.name, cls.model.name,
@ -100,7 +100,7 @@ class TenantService(CommonService):
cls.model.parser_ids, cls.model.parser_ids,
UserTenant.role] UserTenant.role]
return list(cls.model.select(*fields) return list(cls.model.select(*fields)
.join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value))) .join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value) & (UserTenant.role == UserTenantRole.OWNER)))
.where(cls.model.status == StatusEnum.VALID.value).dicts()) .where(cls.model.status == StatusEnum.VALID.value).dicts())
@classmethod @classmethod
@ -115,7 +115,7 @@ class TenantService(CommonService):
cls.model.img2txt_id, cls.model.img2txt_id,
UserTenant.role] UserTenant.role]
return list(cls.model.select(*fields) return list(cls.model.select(*fields)
.join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value) & (UserTenant.role == UserTenantRole.NORMAL.value))) .join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value) & (UserTenant.role == UserTenantRole.NORMAL)))
.where(cls.model.status == StatusEnum.VALID.value).dicts()) .where(cls.model.status == StatusEnum.VALID.value).dicts())
@classmethod @classmethod
@ -143,9 +143,8 @@ class UserTenantService(CommonService):
def get_by_tenant_id(cls, tenant_id): def get_by_tenant_id(cls, tenant_id):
fields = [ fields = [
cls.model.user_id, cls.model.user_id,
cls.model.tenant_id,
cls.model.role,
cls.model.status, cls.model.status,
cls.model.role,
User.nickname, User.nickname,
User.email, User.email,
User.avatar, User.avatar,
@ -153,8 +152,24 @@ class UserTenantService(CommonService):
User.is_active, User.is_active,
User.is_anonymous, User.is_anonymous,
User.status, User.status,
User.update_date,
User.is_superuser] User.is_superuser]
return list(cls.model.select(*fields) return list(cls.model.select(*fields)
.join(User, on=((cls.model.user_id == User.id) & (cls.model.status == StatusEnum.VALID.value))) .join(User, on=((cls.model.user_id == User.id) & (cls.model.status == StatusEnum.VALID.value) & (cls.model.role != UserTenantRole.OWNER)))
.where(cls.model.tenant_id == tenant_id) .where(cls.model.tenant_id == tenant_id)
.dicts()) .dicts())
@classmethod
@DB.connection_context()
def get_tenants_by_user_id(cls, user_id):
fields = [
cls.model.tenant_id,
cls.model.role,
User.nickname,
User.email,
User.avatar,
User.update_date
]
return list(cls.model.select(*fields)
.join(User, on=((cls.model.tenant_id == User.id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value)))
.where(cls.model.status == StatusEnum.VALID.value).dicts())

View File

@ -38,7 +38,7 @@ from api.versions import get_versions
def update_progress(): def update_progress():
while True: while True:
time.sleep(1) time.sleep(3)
try: try:
DocumentService.update_progress() DocumentService.update_progress()
except Exception as e: except Exception as e:

View File

@ -344,3 +344,8 @@ def download_img(url):
return "data:" + \ return "data:" + \
response.headers.get('Content-Type', 'image/jpg') + ";" + \ response.headers.get('Content-Type', 'image/jpg') + ";" + \
"base64," + base64.b64encode(response.content).decode("utf-8") "base64," + base64.b64encode(response.content).decode("utf-8")
def delta_seconds(date_string: str):
dt = datetime.datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S")
return (datetime.datetime.now() - dt).total_seconds()

View File

@ -49,6 +49,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
layers = [(0, len(chunks))] layers = [(0, len(chunks))]
start, end = 0, len(chunks) start, end = 0, len(chunks)
if len(chunks) <= 1: return if len(chunks) <= 1: return
chunks = [(s, a) for s, a in chunks if len(a) > 0]
def summarize(ck_idx, lock): def summarize(ck_idx, lock):
nonlocal chunks nonlocal chunks
@ -64,6 +65,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
print("SUM:", cnt) print("SUM:", cnt)
embds, _ = self._embd_model.encode([cnt]) embds, _ = self._embd_model.encode([cnt])
with lock: with lock:
if not len(embds[0]): return
chunks.append((cnt, embds[0])) chunks.append((cnt, embds[0]))
except Exception as e: except Exception as e:
print(e, flush=True) print(e, flush=True)