KevinHuSh
commited on
Commit
·
98d3e21
1
Parent(s):
11f4a85
refine file select code (#1021)
Browse files### What problem does this PR solve?
#1015
### Type of change
- [x] Refactoring
- api/db/init_data.py +1 -1
- api/db/services/file_service.py +5 -6
api/db/init_data.py
CHANGED
@@ -385,7 +385,7 @@ def init_llm_factory():
|
|
385 |
{
|
386 |
"fid": factory_infos[7]["name"],
|
387 |
"llm_name": "maidalun1020/bce-reranker-base_v1",
|
388 |
-
"tags": "RE-RANK,
|
389 |
"max_tokens": 512,
|
390 |
"model_type": LLMType.RERANK.value
|
391 |
},
|
|
|
385 |
{
|
386 |
"fid": factory_infos[7]["name"],
|
387 |
"llm_name": "maidalun1020/bce-reranker-base_v1",
|
388 |
+
"tags": "RE-RANK, 512",
|
389 |
"max_tokens": 512,
|
390 |
"model_type": LLMType.RERANK.value
|
391 |
},
|
api/db/services/file_service.py
CHANGED
@@ -164,12 +164,11 @@ class FileService(CommonService):
|
|
164 |
@classmethod
|
165 |
@DB.connection_context()
|
166 |
def get_kb_folder(cls, tenant_id):
|
167 |
-
for root in cls.model.select().where(
|
168 |
-
|
169 |
-
for folder in cls.model.select().where(
|
170 |
-
|
171 |
-
|
172 |
-
):
|
173 |
return folder.to_dict()
|
174 |
assert False, "Can't find the KB folder. Database init error."
|
175 |
|
|
|
164 |
@classmethod
|
165 |
@DB.connection_context()
|
166 |
def get_kb_folder(cls, tenant_id):
|
167 |
+
for root in cls.model.select().where(
|
168 |
+
(cls.model.tenant_id == tenant_id), (cls.model.parent_id == cls.model.id)):
|
169 |
+
for folder in cls.model.select().where(
|
170 |
+
(cls.model.tenant_id == tenant_id), (cls.model.parent_id == root.id),
|
171 |
+
(cls.model.name == KNOWLEDGEBASE_FOLDER_NAME)):
|
|
|
172 |
return folder.to_dict()
|
173 |
assert False, "Can't find the KB folder. Database init error."
|
174 |
|