Kevin Hu commited on
Commit
ab87187
·
1 Parent(s): f850783

add API for tenant function (#2866)

Browse files

### 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)

api/apps/canvas_app.py CHANGED
@@ -18,8 +18,6 @@ from functools import partial
18
  from flask import request, Response
19
  from flask_login import login_required, current_user
20
  from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
21
- from api.db.services.dialog_service import full_question
22
- from api.db.services.user_service import TenantService
23
  from api.settings import RetCode
24
  from api.utils import get_uuid
25
  from api.utils.api_utils import get_json_result, server_error_response, validate_request, get_data_error_result
@@ -111,8 +109,9 @@ def run():
111
  if "message" in req:
112
  canvas.messages.append({"role": "user", "content": req["message"], "id": message_id})
113
  if len([m for m in canvas.messages if m["role"] == "user"]) > 1:
114
- ten = TenantService.get_by_user_id(current_user.id)[0]
115
  #req["message"] = full_question(ten["tenant_id"], ten["llm_id"], canvas.messages)
 
116
  canvas.add_user_input(req["message"])
117
  answer = canvas.run(stream=stream)
118
  print(canvas)
 
18
  from flask import request, Response
19
  from flask_login import login_required, current_user
20
  from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
 
 
21
  from api.settings import RetCode
22
  from api.utils import get_uuid
23
  from api.utils.api_utils import get_json_result, server_error_response, validate_request, get_data_error_result
 
109
  if "message" in req:
110
  canvas.messages.append({"role": "user", "content": req["message"], "id": message_id})
111
  if len([m for m in canvas.messages if m["role"] == "user"]) > 1:
112
+ #ten = TenantService.get_info_by(current_user.id)[0]
113
  #req["message"] = full_question(ten["tenant_id"], ten["llm_id"], canvas.messages)
114
+ pass
115
  canvas.add_user_input(req["message"])
116
  answer = canvas.run(stream=stream)
117
  print(canvas)
api/apps/conversation_app.py CHANGED
@@ -218,7 +218,7 @@ def tts():
218
  req = request.json
219
  text = req["text"]
220
 
221
- tenants = TenantService.get_by_user_id(current_user.id)
222
  if not tenants:
223
  return get_data_error_result(retmsg="Tenant not found!")
224
 
 
218
  req = request.json
219
  text = req["text"]
220
 
221
+ tenants = TenantService.get_info_by(current_user.id)
222
  if not tenants:
223
  return get_data_error_result(retmsg="Tenant not found!")
224
 
api/apps/system_app.py CHANGED
@@ -133,15 +133,9 @@ def token_list():
133
  return server_error_response(e)
134
 
135
 
136
- @manager.route('/rm', methods=['POST'])
137
- @validate_request("tokens", "tenant_id")
138
  @login_required
139
- def rm():
140
- req = request.json
141
- try:
142
- for token in req["tokens"]:
143
- APITokenService.filter_delete(
144
- [APIToken.tenant_id == req["tenant_id"], APIToken.token == token])
145
- return get_json_result(data=True)
146
- except Exception as e:
147
- return server_error_response(e)
 
133
  return server_error_response(e)
134
 
135
 
136
+ @manager.route('/token/<token>', methods=['DELETE'])
 
137
  @login_required
138
+ def rm(token):
139
+ APITokenService.filter_delete(
140
+ [APIToken.tenant_id == current_user.id, APIToken.token == token])
141
+ return get_json_result(data=True)
 
 
 
 
 
api/apps/tenant_app.py CHANGED
@@ -15,25 +15,14 @@
15
  #
16
 
17
  from flask import request
18
- from flask_login import current_user, login_required
19
 
20
  from api.db import UserTenantRole, StatusEnum
21
  from api.db.db_models import UserTenant
22
- from api.db.services.user_service import TenantService, UserTenantService
23
- from api.settings import RetCode
24
 
25
- from api.utils import get_uuid
26
- from api.utils.api_utils import get_json_result, validate_request, server_error_response
27
-
28
-
29
- @manager.route("/list", methods=["GET"])
30
- @login_required
31
- def tenant_list():
32
- try:
33
- tenants = TenantService.get_by_user_id(current_user.id)
34
- return get_json_result(data=tenants)
35
- except Exception as e:
36
- return server_error_response(e)
37
 
38
 
39
  @manager.route("/<tenant_id>/user/list", methods=["GET"])
@@ -41,6 +30,8 @@ def tenant_list():
41
  def user_list(tenant_id):
42
  try:
43
  users = UserTenantService.get_by_tenant_id(tenant_id)
 
 
44
  return get_json_result(data=users)
45
  except Exception as e:
46
  return server_error_response(e)
@@ -48,30 +39,31 @@ def user_list(tenant_id):
48
 
49
  @manager.route('/<tenant_id>/user', methods=['POST'])
50
  @login_required
51
- @validate_request("user_id")
52
  def create(tenant_id):
53
- user_id = request.json.get("user_id")
54
- if not user_id:
55
- return get_json_result(
56
- data=False, retmsg='Lack of "USER ID"', retcode=RetCode.ARGUMENT_ERROR)
57
 
58
- try:
59
- user_tenants = UserTenantService.query(user_id=user_id, tenant_id=tenant_id)
60
- if user_tenants:
61
- uuid = user_tenants[0].id
62
- return get_json_result(data={"id": uuid})
63
-
64
- uuid = get_uuid()
65
- UserTenantService.save(
66
- id = uuid,
67
- user_id = user_id,
68
- tenant_id = tenant_id,
69
- role = UserTenantRole.NORMAL.value,
70
- status = StatusEnum.VALID.value)
71
-
72
- return get_json_result(data={"id": uuid})
73
- except Exception as e:
74
- return server_error_response(e)
 
75
 
76
 
77
  @manager.route('/<tenant_id>/user/<user_id>', methods=['DELETE'])
@@ -82,4 +74,25 @@ def rm(tenant_id, user_id):
82
  return get_json_result(data=True)
83
  except Exception as e:
84
  return server_error_response(e)
85
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  #
16
 
17
  from flask import request
18
+ from flask_login import login_required, current_user
19
 
20
  from api.db import UserTenantRole, StatusEnum
21
  from api.db.db_models import UserTenant
22
+ from api.db.services.user_service import UserTenantService, UserService
 
23
 
24
+ from api.utils import get_uuid, delta_seconds
25
+ from api.utils.api_utils import get_json_result, validate_request, server_error_response, get_data_error_result
 
 
 
 
 
 
 
 
 
 
26
 
27
 
28
  @manager.route("/<tenant_id>/user/list", methods=["GET"])
 
30
  def user_list(tenant_id):
31
  try:
32
  users = UserTenantService.get_by_tenant_id(tenant_id)
33
+ for u in users:
34
+ u["delta_seconds"] = delta_seconds(u["update_date"])
35
  return get_json_result(data=users)
36
  except Exception as e:
37
  return server_error_response(e)
 
39
 
40
  @manager.route('/<tenant_id>/user', methods=['POST'])
41
  @login_required
42
+ @validate_request("email")
43
  def create(tenant_id):
44
+ req = request.json
45
+ usrs = UserService.query(email=req["email"])
46
+ if not usrs:
47
+ return get_data_error_result(retmsg="User not found.")
48
 
49
+ user_id = usrs[0].id
50
+ user_tenants = UserTenantService.query(user_id=user_id, tenant_id=tenant_id)
51
+ if user_tenants:
52
+ if user_tenants[0].status == UserTenantRole.NORMAL.value:
53
+ return get_data_error_result(retmsg="This user is in the team already.")
54
+ return get_data_error_result(retmsg="Invitation notification is sent.")
55
+
56
+ UserTenantService.save(
57
+ id=get_uuid(),
58
+ user_id=user_id,
59
+ tenant_id=tenant_id,
60
+ role=UserTenantRole.INVITE,
61
+ status=StatusEnum.VALID.value)
62
+
63
+ usr = list(usrs.dicts())[0]
64
+ usr = {k: v for k, v in usr.items() if k in ["id", "avatar", "email", "nickname"]}
65
+
66
+ return get_json_result(data=usr)
67
 
68
 
69
  @manager.route('/<tenant_id>/user/<user_id>', methods=['DELETE'])
 
74
  return get_json_result(data=True)
75
  except Exception as e:
76
  return server_error_response(e)
77
+
78
+
79
+ @manager.route("/list", methods=["GET"])
80
+ @login_required
81
+ def tenant_list():
82
+ try:
83
+ users = UserTenantService.get_tenants_by_user_id(current_user.id)
84
+ for u in users:
85
+ u["delta_seconds"] = delta_seconds(u["update_date"])
86
+ return get_json_result(data=users)
87
+ except Exception as e:
88
+ return server_error_response(e)
89
+
90
+
91
+ @manager.route("/agree/<tenant_id>", methods=["GET"])
92
+ @login_required
93
+ def agree(tenant_id):
94
+ try:
95
+ UserTenantService.filter_update([UserTenant.tenant_id == tenant_id, UserTenant.user_id == current_user.id], {"role": UserTenantRole.NORMAL})
96
+ return get_json_result(data=True)
97
+ except Exception as e:
98
+ return server_error_response(e)
api/apps/user_app.py CHANGED
@@ -23,7 +23,7 @@ from flask_login import login_required, current_user, login_user, logout_user
23
 
24
  from api.db.db_models import TenantLLM
25
  from api.db.services.llm_service import TenantLLMService, LLMService
26
- from api.utils.api_utils import server_error_response, validate_request
27
  from api.utils import get_uuid, get_format_time, decrypt, download_img, current_timestamp, datetime_format
28
  from api.db import UserTenantRole, LLMType, FileType
29
  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():
402
  @login_required
403
  def tenant_info():
404
  try:
405
- tenants = TenantService.get_by_user_id(current_user.id)[0]
406
- return get_json_result(data=tenants)
 
 
407
  except Exception as e:
408
  return server_error_response(e)
409
 
 
23
 
24
  from api.db.db_models import TenantLLM
25
  from api.db.services.llm_service import TenantLLMService, LLMService
26
+ from api.utils.api_utils import server_error_response, validate_request, get_data_error_result
27
  from api.utils import get_uuid, get_format_time, decrypt, download_img, current_timestamp, datetime_format
28
  from api.db import UserTenantRole, LLMType, FileType
29
  from api.settings import RetCode, GITHUB_OAUTH, FEISHU_OAUTH, CHAT_MDL, EMBEDDING_MDL, ASR_MDL, IMAGE2TEXT_MDL, PARSERS, \
 
402
  @login_required
403
  def tenant_info():
404
  try:
405
+ tenants = TenantService.get_info_by(current_user.id)
406
+ if not tenants:
407
+ return get_data_error_result(retmsg="Tenant not found!")
408
+ return get_json_result(data=tenants[0])
409
  except Exception as e:
410
  return server_error_response(e)
411
 
api/db/__init__.py CHANGED
@@ -27,6 +27,7 @@ class UserTenantRole(StrEnum):
27
  OWNER = 'owner'
28
  ADMIN = 'admin'
29
  NORMAL = 'normal'
 
30
 
31
 
32
  class TenantPermission(StrEnum):
 
27
  OWNER = 'owner'
28
  ADMIN = 'admin'
29
  NORMAL = 'normal'
30
+ INVITE = 'invite'
31
 
32
 
33
  class TenantPermission(StrEnum):
api/db/services/document_service.py CHANGED
@@ -362,7 +362,7 @@ class DocumentService(CommonService):
362
  elif finished:
363
  if d["parser_config"].get("raptor", {}).get("use_raptor") and d["progress_msg"].lower().find(" raptor")<0:
364
  queue_raptor_tasks(d)
365
- prg *= 0.98
366
  msg.append("------ RAPTOR -------")
367
  else:
368
  status = TaskStatus.DONE.value
@@ -379,7 +379,8 @@ class DocumentService(CommonService):
379
  info["progress_msg"] = msg
380
  cls.update_by_id(d["id"], info)
381
  except Exception as e:
382
- stat_logger.error("fetch task exception:" + str(e))
 
383
 
384
  @classmethod
385
  @DB.connection_context()
 
362
  elif finished:
363
  if d["parser_config"].get("raptor", {}).get("use_raptor") and d["progress_msg"].lower().find(" raptor")<0:
364
  queue_raptor_tasks(d)
365
+ prg = 0.98 * len(tsks)/(len(tsks)+1)
366
  msg.append("------ RAPTOR -------")
367
  else:
368
  status = TaskStatus.DONE.value
 
379
  info["progress_msg"] = msg
380
  cls.update_by_id(d["id"], info)
381
  except Exception as e:
382
+ if str(e).find("'0'") < 0:
383
+ stat_logger.error("fetch task exception:" + str(e))
384
 
385
  @classmethod
386
  @DB.connection_context()
api/db/services/user_service.py CHANGED
@@ -87,7 +87,7 @@ class TenantService(CommonService):
87
 
88
  @classmethod
89
  @DB.connection_context()
90
- def get_by_user_id(cls, user_id):
91
  fields = [
92
  cls.model.id.alias("tenant_id"),
93
  cls.model.name,
@@ -100,7 +100,7 @@ class TenantService(CommonService):
100
  cls.model.parser_ids,
101
  UserTenant.role]
102
  return list(cls.model.select(*fields)
103
- .join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value)))
104
  .where(cls.model.status == StatusEnum.VALID.value).dicts())
105
 
106
  @classmethod
@@ -115,7 +115,7 @@ class TenantService(CommonService):
115
  cls.model.img2txt_id,
116
  UserTenant.role]
117
  return list(cls.model.select(*fields)
118
- .join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value) & (UserTenant.role == UserTenantRole.NORMAL.value)))
119
  .where(cls.model.status == StatusEnum.VALID.value).dicts())
120
 
121
  @classmethod
@@ -143,9 +143,8 @@ class UserTenantService(CommonService):
143
  def get_by_tenant_id(cls, tenant_id):
144
  fields = [
145
  cls.model.user_id,
146
- cls.model.tenant_id,
147
- cls.model.role,
148
  cls.model.status,
 
149
  User.nickname,
150
  User.email,
151
  User.avatar,
@@ -153,8 +152,24 @@ class UserTenantService(CommonService):
153
  User.is_active,
154
  User.is_anonymous,
155
  User.status,
 
156
  User.is_superuser]
157
  return list(cls.model.select(*fields)
158
- .join(User, on=((cls.model.user_id == User.id) & (cls.model.status == StatusEnum.VALID.value)))
159
  .where(cls.model.tenant_id == tenant_id)
160
- .dicts())
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
 
88
  @classmethod
89
  @DB.connection_context()
90
+ def get_info_by(cls, user_id):
91
  fields = [
92
  cls.model.id.alias("tenant_id"),
93
  cls.model.name,
 
100
  cls.model.parser_ids,
101
  UserTenant.role]
102
  return list(cls.model.select(*fields)
103
+ .join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value) & (UserTenant.role == UserTenantRole.OWNER)))
104
  .where(cls.model.status == StatusEnum.VALID.value).dicts())
105
 
106
  @classmethod
 
115
  cls.model.img2txt_id,
116
  UserTenant.role]
117
  return list(cls.model.select(*fields)
118
+ .join(UserTenant, on=((cls.model.id == UserTenant.tenant_id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value) & (UserTenant.role == UserTenantRole.NORMAL)))
119
  .where(cls.model.status == StatusEnum.VALID.value).dicts())
120
 
121
  @classmethod
 
143
  def get_by_tenant_id(cls, tenant_id):
144
  fields = [
145
  cls.model.user_id,
 
 
146
  cls.model.status,
147
+ cls.model.role,
148
  User.nickname,
149
  User.email,
150
  User.avatar,
 
152
  User.is_active,
153
  User.is_anonymous,
154
  User.status,
155
+ User.update_date,
156
  User.is_superuser]
157
  return list(cls.model.select(*fields)
158
+ .join(User, on=((cls.model.user_id == User.id) & (cls.model.status == StatusEnum.VALID.value) & (cls.model.role != UserTenantRole.OWNER)))
159
  .where(cls.model.tenant_id == tenant_id)
160
+ .dicts())
161
+
162
+ @classmethod
163
+ @DB.connection_context()
164
+ def get_tenants_by_user_id(cls, user_id):
165
+ fields = [
166
+ cls.model.tenant_id,
167
+ cls.model.role,
168
+ User.nickname,
169
+ User.email,
170
+ User.avatar,
171
+ User.update_date
172
+ ]
173
+ return list(cls.model.select(*fields)
174
+ .join(User, on=((cls.model.tenant_id == User.id) & (UserTenant.user_id == user_id) & (UserTenant.status == StatusEnum.VALID.value)))
175
+ .where(cls.model.status == StatusEnum.VALID.value).dicts())
api/ragflow_server.py CHANGED
@@ -38,7 +38,7 @@ from api.versions import get_versions
38
 
39
  def update_progress():
40
  while True:
41
- time.sleep(1)
42
  try:
43
  DocumentService.update_progress()
44
  except Exception as e:
 
38
 
39
  def update_progress():
40
  while True:
41
+ time.sleep(3)
42
  try:
43
  DocumentService.update_progress()
44
  except Exception as e:
api/utils/__init__.py CHANGED
@@ -344,3 +344,8 @@ def download_img(url):
344
  return "data:" + \
345
  response.headers.get('Content-Type', 'image/jpg') + ";" + \
346
  "base64," + base64.b64encode(response.content).decode("utf-8")
 
 
 
 
 
 
344
  return "data:" + \
345
  response.headers.get('Content-Type', 'image/jpg') + ";" + \
346
  "base64," + base64.b64encode(response.content).decode("utf-8")
347
+
348
+
349
+ def delta_seconds(date_string: str):
350
+ dt = datetime.datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S")
351
+ return (datetime.datetime.now() - dt).total_seconds()
rag/raptor.py CHANGED
@@ -49,6 +49,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
49
  layers = [(0, len(chunks))]
50
  start, end = 0, len(chunks)
51
  if len(chunks) <= 1: return
 
52
 
53
  def summarize(ck_idx, lock):
54
  nonlocal chunks
@@ -64,6 +65,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
64
  print("SUM:", cnt)
65
  embds, _ = self._embd_model.encode([cnt])
66
  with lock:
 
67
  chunks.append((cnt, embds[0]))
68
  except Exception as e:
69
  print(e, flush=True)
 
49
  layers = [(0, len(chunks))]
50
  start, end = 0, len(chunks)
51
  if len(chunks) <= 1: return
52
+ chunks = [(s, a) for s, a in chunks if len(a) > 0]
53
 
54
  def summarize(ck_idx, lock):
55
  nonlocal chunks
 
65
  print("SUM:", cnt)
66
  embds, _ = self._embd_model.encode([cnt])
67
  with lock:
68
+ if not len(embds[0]): return
69
  chunks.append((cnt, embds[0]))
70
  except Exception as e:
71
  print(e, flush=True)