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 +2 -3
- api/apps/conversation_app.py +1 -1
- api/apps/system_app.py +5 -11
- api/apps/tenant_app.py +51 -38
- api/apps/user_app.py +5 -3
- api/db/__init__.py +1 -0
- api/db/services/document_service.py +3 -2
- api/db/services/user_service.py +22 -7
- api/ragflow_server.py +1 -1
- api/utils/__init__.py +5 -0
- rag/raptor.py +2 -0
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.
|
| 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.
|
| 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('/
|
| 137 |
-
@validate_request("tokens", "tenant_id")
|
| 138 |
@login_required
|
| 139 |
-
def rm():
|
| 140 |
-
|
| 141 |
-
|
| 142 |
-
|
| 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
|
| 19 |
|
| 20 |
from api.db import UserTenantRole, StatusEnum
|
| 21 |
from api.db.db_models import UserTenant
|
| 22 |
-
from api.db.services.user_service import
|
| 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("
|
| 52 |
def create(tenant_id):
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
|
| 57 |
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
return
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
|
|
|
| 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.
|
| 406 |
-
|
|
|
|
|
|
|
| 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
|
| 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 |
-
|
|
|
|
| 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
|
| 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
|
| 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(
|
| 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)
|