KevinHuSh
commited on
Commit
·
1d694ff
1
Parent(s):
66d269d
correct mismatched kb doc number (#826)
Browse files### What problem does this PR solve?
#620
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
api/db/init_data.py
CHANGED
@@ -21,6 +21,8 @@ from copy import deepcopy
|
|
21 |
from api.db import LLMType, UserTenantRole
|
22 |
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM
|
23 |
from api.db.services import UserService
|
|
|
|
|
24 |
from api.db.services.llm_service import LLMFactoriesService, LLMService, TenantLLMService, LLMBundle
|
25 |
from api.db.services.user_service import TenantService, UserTenantService
|
26 |
from api.settings import CHAT_MDL, EMBEDDING_MDL, ASR_MDL, IMAGE2TEXT_MDL, PARSERS, LLM_FACTORY, API_KEY, LLM_BASE_URL
|
@@ -406,6 +408,8 @@ def init_llm_factory():
|
|
406 |
except Exception as e:
|
407 |
pass
|
408 |
break
|
|
|
|
|
409 |
"""
|
410 |
drop table llm;
|
411 |
drop table llm_factories;
|
|
|
21 |
from api.db import LLMType, UserTenantRole
|
22 |
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM
|
23 |
from api.db.services import UserService
|
24 |
+
from api.db.services.document_service import DocumentService
|
25 |
+
from api.db.services.knowledgebase_service import KnowledgebaseService
|
26 |
from api.db.services.llm_service import LLMFactoriesService, LLMService, TenantLLMService, LLMBundle
|
27 |
from api.db.services.user_service import TenantService, UserTenantService
|
28 |
from api.settings import CHAT_MDL, EMBEDDING_MDL, ASR_MDL, IMAGE2TEXT_MDL, PARSERS, LLM_FACTORY, API_KEY, LLM_BASE_URL
|
|
|
408 |
except Exception as e:
|
409 |
pass
|
410 |
break
|
411 |
+
for kb_id in KnowledgebaseService.get_all_ids():
|
412 |
+
KnowledgebaseService.update_by_id(kb_id, {"doc_num": DocumentService.get_kb_doc_count(kb_id)})
|
413 |
"""
|
414 |
drop table llm;
|
415 |
drop table llm_factories;
|
api/db/services/document_service.py
CHANGED
@@ -265,3 +265,9 @@ class DocumentService(CommonService):
|
|
265 |
except Exception as e:
|
266 |
stat_logger.error("fetch task exception:" + str(e))
|
267 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
265 |
except Exception as e:
|
266 |
stat_logger.error("fetch task exception:" + str(e))
|
267 |
|
268 |
+
@classmethod
|
269 |
+
@DB.connection_context()
|
270 |
+
def get_kb_doc_count(cls, kb_id):
|
271 |
+
return len(cls.model.select(cls.model.id).where(
|
272 |
+
cls.model.kb_id == kb_id).dicts())
|
273 |
+
|
api/db/services/knowledgebase_service.py
CHANGED
@@ -112,3 +112,8 @@ class KnowledgebaseService(CommonService):
|
|
112 |
if kb:
|
113 |
return True, kb[0]
|
114 |
return False, None
|
|
|
|
|
|
|
|
|
|
|
|
112 |
if kb:
|
113 |
return True, kb[0]
|
114 |
return False, None
|
115 |
+
|
116 |
+
@classmethod
|
117 |
+
@DB.connection_context()
|
118 |
+
def get_all_ids(cls):
|
119 |
+
return [m["id"] for m in cls.model.select(cls.model.id).dicts()]
|