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()]
|