Langchain-Chatchat

1
from server.db.models.knowledge_base_model import KnowledgeBaseModel
2
from server.db.session import with_session
3

4

5
@with_session
6
def add_kb_to_db(session, kb_name, kb_info, vs_type, embed_model):
7
    # 创建知识库实例
8
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
9
    if not kb:
10
        kb = KnowledgeBaseModel(kb_name=kb_name, kb_info=kb_info, vs_type=vs_type, embed_model=embed_model)
11
        session.add(kb)
12
    else:  # update kb with new vs_type and embed_model
13
        kb.kb_info = kb_info
14
        kb.vs_type = vs_type
15
        kb.embed_model = embed_model
16
    return True
17

18

19
@with_session
20
def list_kbs_from_db(session, min_file_count: int = -1):
21
    kbs = session.query(KnowledgeBaseModel.kb_name).filter(KnowledgeBaseModel.file_count > min_file_count).all()
22
    kbs = [kb[0] for kb in kbs]
23
    return kbs
24

25

26
@with_session
27
def kb_exists(session, kb_name):
28
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
29
    status = True if kb else False
30
    return status
31

32

33
@with_session
34
def load_kb_from_db(session, kb_name):
35
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
36
    if kb:
37
        kb_name, vs_type, embed_model = kb.kb_name, kb.vs_type, kb.embed_model
38
    else:
39
        kb_name, vs_type, embed_model = None, None, None
40
    return kb_name, vs_type, embed_model
41

42

43
@with_session
44
def delete_kb_from_db(session, kb_name):
45
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
46
    if kb:
47
        session.delete(kb)
48
    return True
49

50

51
@with_session
52
def get_kb_detail(session, kb_name: str) -> dict:
53
    kb: KnowledgeBaseModel = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
54
    if kb:
55
        return {
56
            "kb_name": kb.kb_name,
57
            "kb_info": kb.kb_info,
58
            "vs_type": kb.vs_type,
59
            "embed_model": kb.embed_model,
60
            "file_count": kb.file_count,
61
            "create_time": kb.create_time,
62
        }
63
    else:
64
        return {}
65