diff --git a/server/db/repository/knowledge_base_repository.py b/server/db/repository/knowledge_base_repository.py index 3f50c3d..cef420a 100644 --- a/server/db/repository/knowledge_base_repository.py +++ b/server/db/repository/knowledge_base_repository.py @@ -42,3 +42,18 @@ def delete_kb_from_db(session, kb_name): if kb: session.delete(kb) return True + + +@with_session +def get_kb_detail(session, kb_name: str) -> dict: + kb: KnowledgeBaseModel = session.query(KnowledgeBaseModel).filter_by(kb_name=kb_name).first() + if kb: + return { + "kb_name": kb.kb_name, + "vs_type": kb.vs_type, + "embed_model": kb.embed_model, + "file_count": kb.file_count, + "create_time": kb.create_time, + } + else: + return {} diff --git a/server/db/repository/knowledge_file_repository.py b/server/db/repository/knowledge_file_repository.py index cdb5c78..9aa00b7 100644 --- a/server/db/repository/knowledge_file_repository.py +++ b/server/db/repository/knowledge_file_repository.py @@ -56,3 +56,22 @@ def doc_exists(session, kb_file: KnowledgeFile): existing_file = session.query(KnowledgeFileModel).filter_by(file_name=kb_file.filename, kb_name=kb_file.kb_name).first() return True if existing_file else False + + +@with_session +def get_file_detail(session, kb_name: str, filename: str) -> dict: + file: KnowledgeFileModel = (session.query(KnowledgeFileModel) + .filter_by(file_name=filename, + kb_name=kb_name).first()) + if file: + return { + "kb_name": file.kb_name, + "file_name": file.file_name, + "file_ext": file.file_ext, + "file_version": file.file_version, + "document_loader": file.document_loader_name, + "text_splitter": file.text_splitter_name, + "create_time": file.create_time, + } + else: + return {}