Langchain-Chatchat

1
from server.db.models.knowledge_metadata_model import SummaryChunkModel
2
from server.db.session import with_session
3
from typing import List, Dict
4

5

6
@with_session
7
def list_summary_from_db(session,
8
                         kb_name: str,
9
                         metadata: Dict = {},
10
                         ) -> List[Dict]:
11
    '''
12
    列出某知识库chunk summary。
13
    返回形式：[{"id": str, "summary_context": str, "doc_ids": str}, ...]
14
    '''
15
    docs = session.query(SummaryChunkModel).filter(SummaryChunkModel.kb_name.ilike(kb_name))
16

17
    for k, v in metadata.items():
18
        docs = docs.filter(SummaryChunkModel.meta_data[k].as_string() == str(v))
19

20
    return [{"id": x.id,
21
             "summary_context": x.summary_context,
22
             "summary_id": x.summary_id,
23
             "doc_ids": x.doc_ids,
24
             "metadata": x.metadata} for x in docs.all()]
25

26

27
@with_session
28
def delete_summary_from_db(session,
29
                           kb_name: str
30
                           ) -> List[Dict]:
31
    '''
32
    删除知识库chunk summary，并返回被删除的Dchunk summary。
33
    返回形式：[{"id": str, "summary_context": str, "doc_ids": str}, ...]
34
    '''
35
    docs = list_summary_from_db(kb_name=kb_name)
36
    query = session.query(SummaryChunkModel).filter(SummaryChunkModel.kb_name.ilike(kb_name))
37
    query.delete(synchronize_session=False)
38
    session.commit()
39
    return docs
40

41

42
@with_session
43
def add_summary_to_db(session,
44
                      kb_name: str,
45
                      summary_infos: List[Dict]):
46
    '''
47
    将总结信息添加到数据库。
48
    summary_infos形式：[{"summary_context": str, "doc_ids": str}, ...]
49
    '''
50
    for summary in summary_infos:
51
        obj = SummaryChunkModel(
52
            kb_name=kb_name,
53
            summary_context=summary["summary_context"],
54
            summary_id=summary["summary_id"],
55
            doc_ids=summary["doc_ids"],
56
            meta_data=summary["metadata"],
57
        )
58
        session.add(obj)
59

60
    session.commit()
61
    return True
62

63

64
@with_session
65
def count_summary_from_db(session, kb_name: str) -> int:
66
    return session.query(SummaryChunkModel).filter(SummaryChunkModel.kb_name.ilike(kb_name)).count()
67