Skip to content

Commit

Permalink
correct mismatched kb doc number (infiniflow#826)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?

infiniflow#620

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
  • Loading branch information
KevinHuSh committed May 17, 2024
1 parent f6a5994 commit a1586e0
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 0 deletions.
4 changes: 4 additions & 0 deletions api/db/init_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
from api.db import LLMType, UserTenantRole
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM
from api.db.services import UserService
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.llm_service import LLMFactoriesService, LLMService, TenantLLMService, LLMBundle
from api.db.services.user_service import TenantService, UserTenantService
from api.settings import CHAT_MDL, EMBEDDING_MDL, ASR_MDL, IMAGE2TEXT_MDL, PARSERS, LLM_FACTORY, API_KEY, LLM_BASE_URL
Expand Down Expand Up @@ -406,6 +408,8 @@ def init_llm_factory():
except Exception as e:
pass
break
for kb_id in KnowledgebaseService.get_all_ids():
KnowledgebaseService.update_by_id(kb_id, {"doc_num": DocumentService.get_kb_doc_count(kb_id)})
"""
drop table llm;
drop table llm_factories;
Expand Down
6 changes: 6 additions & 0 deletions api/db/services/document_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,3 +265,9 @@ def update_progress(cls):
except Exception as e:
stat_logger.error("fetch task exception:" + str(e))

@classmethod
@DB.connection_context()
def get_kb_doc_count(cls, kb_id):
return len(cls.model.select(cls.model.id).where(
cls.model.kb_id == kb_id).dicts())

5 changes: 5 additions & 0 deletions api/db/services/knowledgebase_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,3 +112,8 @@ def get_by_name(cls, kb_name, tenant_id):
if kb:
return True, kb[0]
return False, None

@classmethod
@DB.connection_context()
def get_all_ids(cls):
return [m["id"] for m in cls.model.select(cls.model.id).dicts()]

0 comments on commit a1586e0

Please sign in to comment.