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