mirror of
https://github.com/RYDE-WORK/Langchain-Chatchat.git
synced 2026-01-19 21:37:20 +08:00
5d422ca9a1
改变 Embeddings 模型改为使用框架 API，不再手动加载，删除自定义 Embeddings Keyword 代码 修改依赖文件，移除 torch transformers 等重依赖 暂时移出对 loom 的集成 后续： 1、优化目录结构 2、检查合并中有无被覆盖的 0.2.10 内容
65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
from server.db.models.knowledge_base_model import KnowledgeBaseModel
|
|
from server.db.session import with_session
|
|
|
|
|
|
@with_session
|
|
def add_kb_to_db(session, kb_name, kb_info, vs_type, embed_model):
|
|
# 创建知识库实例
|
|
kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
|
|
if not kb:
|
|
kb = KnowledgeBaseModel(kb_name=kb_name, kb_info=kb_info, vs_type=vs_type, embed_model=embed_model)
|
|
session.add(kb)
|
|
else: # update kb with new vs_type and embed_model
|
|
kb.kb_info = kb_info
|
|
kb.vs_type = vs_type
|
|
kb.embed_model = embed_model
|
|
return True
|
|
|
|
|
|
@with_session
|
|
def list_kbs_from_db(session, min_file_count: int = -1):
|
|
kbs = session.query(KnowledgeBaseModel.kb_name).filter(KnowledgeBaseModel.file_count > min_file_count).all()
|
|
kbs = [kb[0] for kb in kbs]
|
|
return kbs
|
|
|
|
|
|
@with_session
|
|
def kb_exists(session, kb_name):
|
|
kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
|
|
status = True if kb else False
|
|
return status
|
|
|
|
|
|
@with_session
|
|
def load_kb_from_db(session, kb_name):
|
|
kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
|
|
if kb:
|
|
kb_name, vs_type, embed_model = kb.kb_name, kb.vs_type, kb.embed_model
|
|
else:
|
|
kb_name, vs_type, embed_model = None, None, None
|
|
return kb_name, vs_type, embed_model
|
|
|
|
|
|
@with_session
|
|
def delete_kb_from_db(session, kb_name):
|
|
kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
|
|
if kb:
|
|
session.delete(kb)
|
|
return True
|
|
|
|
|
|
@with_session
|
|
def get_kb_detail(session, kb_name: str) -> dict:
|
|
kb: KnowledgeBaseModel = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
|
|
if kb:
|
|
return {
|
|
"kb_name": kb.kb_name,
|
|
"kb_info": kb.kb_info,
|
|
"vs_type": kb.vs_type,
|
|
"embed_model": kb.embed_model,
|
|
"file_count": kb.file_count,
|
|
"create_time": kb.create_time,
|
|
}
|
|
else:
|
|
return {}
|