Lightrag1.0.3代码解读

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

# ====== 文章处理 ======
# 单个字符串也变成List(str)
if isinstance(string_or_strings, str):
    string_or_strings = [string_or_strings]

# {hash("doc-"+txt_i的字符串):{"content": txt_i.strip()}}，txt_i是第i个文章
new_docs = {
    compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
    for c in string_or_strings
}

'''
n个 hash("doc-"+txt_i的字符串) 组成的list 传入 self.full_docs.filter_keys
self.full_docs.filter_keys 是 JsonKVStorage.filter_keys方法
返回set{str}(以字符串为元素的字典)
作用是：某文章 和历史数据中重复 或 和传入数据中其他文章重复 则排除。
通过比对hash值
'''
_add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))
# key在_add_doc_keys里的kv对留下
new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}

# 日志登记
if not len(new_docs):
    logger.warning("All docs are already in the storage")
    return
# storage 是 要更新的flag
update_storage = True
logger.info(f"[New Docs] inserting {len(new_docs)} docs")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

# ====== 文本块处理 ======

# 一次ainsert中的多个文章的chunks字典汇总，ainsert只是事件循环中的一个事件
# {"chunk-"+chunk字符串哈希值: dict, 哈希值: dict,...}
inserting_chunks = {}

# 外循环 第i篇文章，内循环，第i篇文章第j个chunk
for doc_key, doc in \
tqdm_async(new_docs.items(), desc="Chunking documents", unit="doc"):

    '''
    inserting_chunks =
    {
        hash("chunk-"+chunk_i的字符串):
                    {"tokens": chunk_i的token数,
                     "content": chunk_i具体内容,
                     "chunk_order_index": 是该文章的第i个chunk,
                     "full_doc_id": doc_key},
        hash("chunk-"+chunk_i+1的字符串):
                    {"tokens": chunk_i+1的token数,
                     "content": chunk_i+1具体内容,
                     "chunk_order_index": 是该文章的第i+1个chunk,
                     "full_doc_id": doc_key},
    }            
    chunks = {哈希值: dict, 哈希值: dict,...}     
    '''

    chunks = {

        # dp是 dict[str, Union[int, str]]，即是个字典，然后键为key，值为Union[int,str]即int 或 str
        # Union[int,str]具体来说，可以是 该chunk的token数，chunk具体内容，是该文章的第i个chunk
        # **dp表示解包字典（去掉大括号），并要求"full_doc_id": doc_key满足dict[str, Union[int, str]]
        # chunking_by_token_size 根据 token 数量切分文本，返回List[dict]吧

        compute_mdhash_id(dp["content"], prefix="chunk-"):
            {**dp, "full_doc_id": doc_key,}

        for dp in chunking_by_token_size(
            doc["content"],
            overlap_token_size=self.chunk_overlap_token_size,
            max_token_size=self.chunk_token_size,
            tiktoken_model=self.tiktoken_model_name,
        )
    }
    
    # dict.update(dict)如果已存在相同的键，则对应的值会被覆盖；如果不存在，则添加新的键值对。
    # dict1 = {'a': 1, 'b': 2}
    # dict2 = {'b': 3, 'c': 4}
    # dict1.update(dict2)
    # print(dict1)  # Output: {'a': 1, 'b': 3, 'c': 4}
    
    # 因为用了hash内容包括chunk内容，所以不重复
    inserting_chunks.update(chunks)

'''
n个 hash("chunk-"+chunk_i的字符串) 组成的list传入self.full_docs.filter_keys
self.full_docs.filter_keys 即 JsonKVStorage.filter_keys
返回set{str}(以字符串为元素的字典)，作用是：某chunk 和历史数据中重复或和传入数据中其他chunk重复则排除。
通过比对hash值
'''
_add_chunk_keys = await self.text_chunks.filter_keys(
    list(inserting_chunks.keys())
)
# 根据key在_add_doc_keys里对inserting_chunks筛选合格的kv对
inserting_chunks = {
    k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
}

# 日志登记
if not len(inserting_chunks):
    logger.warning("All chunks are already in the storage")
    return
logger.info(f"[New Chunks] inserting {len(inserting_chunks)} chunks")

# 将inserting_chunks(字典)存入json
await self.chunks_vdb.upsert(inserting_chunks)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

# ====== 实体抽取 ======
logger.info("[Entity Extraction]...")

# asdict来自“from dataclasses import asdict“，能构建Lightrag类的属性字典
# maybe_new_kg为更新过后的knowledge_graph_inst，是NetworkXStorage类的实例
# maybe_new_kg最重要的属性为 _graph，其为nx.Graph()
maybe_new_kg = await extract_entities(
    inserting_chunks,
    knowledge_graph_inst=self.chunk_entity_relation_graph,
    entity_vdb=self.entities_vdb,
    relationships_vdb=self.relationships_vdb,
    global_config=asdict(self),
)

# operate.py中 if not len(all_entities_data) if not len(all_relationships_data)会return None
if maybe_new_kg is None:
    logger.warning("No new entities and relationships found")
    return

self.chunk_entity_relation_graph = maybe_new_kg

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246

async def extract_entities(
    chunks: dict[str, TextChunkSchema],
    knowledge_graph_inst: BaseGraphStorage,
    entity_vdb: BaseVectorStorage,
    relationships_vdb: BaseVectorStorage,
    global_config: dict,
) -> Union[BaseGraphStorage, None]:

    '''
    global_config["llm_model_func"]
     -> gpt_4o_mini_complete(prompt, system_prompt=None, history_messages=[], **kwargs)
     -> openai_complete_if_cache(model, prompt, system_prompt=None, history_messages=[],
     base_url=None, api_key=None, **kwargs)  后来的封装了，可以选择模型，不像这里函数名写明了4o_min
    上一行三个参数没有传入和调用实际上
     '''
    use_llm_func: callable = global_config["llm_model_func"]
    # 默认设置为1，貌似和多轮对话有关
    entity_extract_max_gleaning = global_config["entity_extract_max_gleaning"]
    '''
        ordered_chunks =
        [
        (hash("chunk-"+chunk_i的字符串),
        {"tokens": chunk_i的token数,
         "content": chunk_i具体内容,
         "chunk_order_index": 是该文章的第i个chunk,
         "full_doc_id": doc_key}),
         (a,b),...
        ]
        即可索引的chunks               
    '''
    ordered_chunks = list(chunks.items())
    '''
    add language and example number params to prompt
    additional parameter 包括 语言 和 例子数量
    实例化Lightrag的时候
    可以添加参数addon_params={"example_number": 1, "language": "Simplfied Chinese"}
    默认英文（PROMPTS["DEFAULT_LANGUAGE"]），None
    '''
    language = global_config["addon_params"].get(
        "language", PROMPTS["DEFAULT_LANGUAGE"]
    )
    example_number = global_config["addon_params"].get("example_number", None)

    # 如果example_number不为None，按设置的个数来，且不超范围，
    # 如果为none，设置为 PROMPTS["entity_extraction_examples"]
    # examples="example1"+"\n"+"example2"+"\n"+"example3"
    if example_number and example_number < len(PROMPTS["entity_extraction_examples"]):
        examples = "\n".join(
            PROMPTS["entity_extraction_examples"][: int(example_number)]
        )
    else:
        examples = "\n".join(PROMPTS["entity_extraction_examples"])

    # tuple_delimiter="<|>" record_delimiter="##" completion_delimiter="<|COMPLETE|>"
    # entity_types=["organization", "person", "geo", "event"]
    example_context_base = dict(
        tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
        record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
        completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
        entity_types=",".join(PROMPTS["DEFAULT_ENTITY_TYPES"]),
        language=language,
    )
    '''
    添加exampel的格式
    
    examples字符串部分内容如下：
    ("实体"{tuple_delimiter}"Jordan"{tuple_delimiter}"人物"{tuple_delimiter}"Jordan 对探索有着共同的承诺，并与 Taylor 就设备进行了重要互动。"){record_delimiter}
    ("关系"{tuple_delimiter}"Jordan"{tuple_delimiter}"Cruz"{tuple_delimiter}"Jordan 对探索的承诺是对 Cruz 控制与秩序愿景的反抗。"{tuple_delimiter}"意识形态冲突, 反抗"{tuple_delimiter}5){record_delimiter}
    ("content_keywords"{tuple_delimiter}"power dynamics, ideological conflict, discovery, rebellion"){completion_delimiter}
    
    examples字符串部分内容说明：
    每行是对chunk的实体抽取、关系抽取或整体内容关键词提取
    开头必为实体、关系或关键词，然后用{tuple_delimiter}隔开，每行记录1个关系过实体，然后行末添加{record_delimiter}
    元数据附加：为关系添加分类标签（如“权力动态”“意识形态冲突”）和强度评分（如 Taylor-Jordan 的关系强度为 8）
    
    examples.format(**dict)将类如{tuple_delimiter}的内容替换为dict中的值（ tuple_delimiter: 值）
    
    最终examples是 {...} 被替换了的一个字符串
    '''
    examples = examples.format(**example_context_base)

    '''
    entity_extract_prompt主要包括：
    -Goal- ...
    -Steps- ...
    -Examples- ...
    -Real Data- ...
    '''
    entity_extract_prompt = PROMPTS["entity_extraction"]

    context_base = dict(
        tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
        record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
        completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
        entity_types=",".join(PROMPTS["DEFAULT_ENTITY_TYPES"]),
        examples=examples,
        language=language,
    )
    # 让大模型继续按上次的格式继续提取
    # """MANY entities were missed in the last extraction.
    # Add them below using the same format:"""
    continue_prompt = PROMPTS["entiti_continue_extraction"]

    # 询问LLM是否遗漏许多 实体 ，回答 是/否
    # """It appears some entities may have still been missed.
    # Answer YES | NO if there are still entities that need to be added."""
    if_loop_prompt = PROMPTS["entiti_if_loop_extraction"]

    already_processed = 0
    already_entities = 0
    already_relations = 0

    # return dict(maybe_nodes), dict(maybe_edges)
    # 即 返回一个 元组，包含两个字典：(nodes_dict, edges_dict)
    # nodes_dict = {str: dict...}
    # edges_dict = {(str, str): dict}
    async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]):...

    # =====提取出每个chunk的entity和relation=====

    # results = [chunk_1的(nodes_dict, edges_dict), chunk_2的(nodes_dict, edges_dict)...]
    #         = [
    #           ( {str: list[dict]} , {(str, str): list[dict]} ),
    #           ( {str: list[dict]} , {(str, str): list[dict]} ),
    #           ...
    #           ]
    # 调用_process_single_content 返回的二元组 即 result
    results = []
    for result in tqdm_async(
        asyncio.as_completed([_process_single_content(c) for c in ordered_chunks]),
        total=len(ordered_chunks),
        desc="Extracting entities from chunks",
        unit="chunk",
    ):
        results.append(await result)

    maybe_nodes = defaultdict(list)
    maybe_edges = defaultdict(list)

    # m_nodes = {str: list[dict]}合并成{str: list[dict,dict,...]}，dict的数量大于等于
    # m_edges = {(str, str): list[dict]}合并同理
    # 将list[二元tuple] results 处理为结点和边
    # 因为results中的各个(nodes, edges)元组是以chunk单位提取的，提取的实体和关系有蛮多重复的
    # 所以这里{str: list[dict,dict,...]} 会对应多个dict
    for m_nodes, m_edges in results:
        # 不同chunk同样实体的dict放在一个list了
        for k, v in m_nodes.items():
            maybe_nodes[k].extend(v)
        # 同理
        for k, v in m_edges.items():
            maybe_edges[tuple(sorted(k))].extend(v)

    # =====添加实体（结点）到图数据结构=====
    logger.info("Inserting entities into storage...")
    all_entities_data = []

    # knowledge_graph_inst为NetworkXStorage实例
    # for result in 列表生成式: all_e....append(await result)
    # 列表生成式=[_merge_upsert.(k, v,...graph_inst...) for k,v in...]

    # 不同chunk提取的实体和关系有蛮多重复的，{str: list[dict,dict,...]}会对应多个dict
    # maybe_nodes：所有chunk同样实体的dict放在一个list了，{str: list[dict,dict,...]}
    # result = node_data = dict(
    #     entity_name=entity_name
    #     entity_type=entity_type,
    #     description=description,
    #     source_id=source_id,
    # )
    # 这里的一个result是集成该实体在 所有的chunk和过往graph里
    # 提取出的dscri、源文章id，出现次数最多的type
    for result in tqdm_async(
        asyncio.as_completed(
            [
                _merge_nodes_then_upsert(k, v, knowledge_graph_inst, global_config)
                for k, v in maybe_nodes.items()
            ]
        ),
        total=len(maybe_nodes),
        desc="Inserting entities",
        unit="entity",
    ):
        all_entities_data.append(await result)

    # =====添加关系（边）到图数据结构=====
    logger.info("Inserting relationships into storage...")
    all_relationships_data = []
    # all_relationships_data = List[dict]
    # result = edge_data = dict(
    #     src_id=src_id,
    #     tgt_id=tgt_id,
    #     description=description,
    #     keywords=keywords,
    # )
    # 多个chunk提取出的dscri用<SEP>连接起来，keywords去重排序后用<SEP>连接
    # k[0]是(node1->str,node2->str), k[1]是[dict, dict, ...]每个dict来自不同chunk
    for result in tqdm_async(
        asyncio.as_completed(
            [
                _merge_edges_then_upsert(
                    k[0], k[1], v, knowledge_graph_inst, global_config
                )
                for k, v in maybe_edges.items()
            ]
        ),
        total=len(maybe_edges),
        desc="Inserting relationships",
        unit="relationship",
    ):
        all_relationships_data.append(await result)

    if not len(all_entities_data):
        logger.warning("Didn't extract any entities, maybe your LLM is not working")
        return None
    if not len(all_relationships_data):
        logger.warning(
            "Didn't extract any relationships, maybe your LLM is not working"
        )
        return None

    # =====存实体向量数据=====
    if entity_vdb is not None:
        data_for_vdb = {
            compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
                "content": dp["entity_name"] + dp["description"],
                "entity_name": dp["entity_name"],
            }
            for dp in all_entities_data
        }
        await entity_vdb.upsert(data_for_vdb)

    # =====存关系向量数据=====
    if relationships_vdb is not None:
        data_for_vdb = {
            compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): {
                "src_id": dp["src_id"],
                "tgt_id": dp["tgt_id"],
                "content": dp["keywords"]
                + dp["src_id"]
                + dp["tgt_id"]
                + dp["description"],
            }
            for dp in all_relationships_data
        }
        await relationships_vdb.upsert(data_for_vdb)

    return knowledge_graph_inst

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210

async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]):
    nonlocal already_processed, already_entities, already_relations

    # chunk_key = hash('chunk-'+chunk_i具体内容)
    chunk_key = chunk_key_dp[0]
    # chunk_dp 见 lightrag.py中的 chunks变量前的 注释，是个字典
    chunk_dp = chunk_key_dp[1]
    # content即chunk_i具体内容
    content = chunk_dp["content"]

    """
    替换初始prompt模板里的{...}
    hint_prompt = entity_extract_prompt.format(**context_base, input_text=content)
    两次.format是因为担心第一次.format后,字典context_base的值中会有{input_text}
    故第一次.format中，用input_text="{input_text}"保留input_text不变，第二次再全部替换
    hint_prompt即终极版prompt
    """

    hint_prompt = entity_extract_prompt.format(
        **context_base, input_text="{input_text}"
    ).format(**context_base, input_text=content)


    # llm抽取实体
    '''
    use_llm_func 即 gpt_4o_mini_complete(...) -> str
    use_llm_func: callable = global_config["llm_model_func"]
    llm_model_func: callable = gpt_4o_mini_complete   --- lightrag.py定义的类属性，global_config
    async def gpt_4o_mini_complete(...)    --- llm.py

    gpt_4o_mini_complete 返回：
    return await openai_complete_if_cache(
          "gpt-4o-mini",
          prompt,
          system_prompt=system_prompt,
          history_messages=history_messages,
          **kwargs,
      )
    final_result是str
    '''
    final_result = await use_llm_func(hint_prompt)

    '''
    一个列表
    history = [
            {"role": 'user', "content": hint_prompt}
            {"role": 'assistant', "content": final_result}
    ]
    '''
    history = pack_user_ass_to_openai_messages(hint_prompt, final_result)

    # "拾穗" 即重复抽取，防止遗漏
    # entity_extract_max_gleaning默认为1
    for now_glean_index in range(entity_extract_max_gleaning):
        # 直接告诉它遗漏了实体，要继续提取
        # """MANY entities were missed in the last extraction.
        # Add them below using the same format:"""
        glean_result = await use_llm_func(continue_prompt, history_messages=history)

        '''
        [
            {"role": 'user', "content": hint_prompt}
            {"role": 'assistant', "content": final_result}
            {"role": 'user', "content": continue_prompt}
            {"role": 'assistant', "content": glean_result}
        ]
        '''
        history += pack_user_ass_to_openai_messages(continue_prompt, glean_result)
        # str+str 字符串相连
        final_result += glean_result

        # 达到设定的最大次数，跳出
        if now_glean_index == entity_extract_max_gleaning - 1:
            break

        #未达到，让LLM判断是否需要再次 拾穗
        if_loop_result: str = await use_llm_func(
            if_loop_prompt, history_messages=history
        )
        if_loop_result = if_loop_result.strip().strip('"').strip("'").lower()
        # 不需要则break
        if if_loop_result != "yes":
            break

    '''
    假如说final result

    ("实体"{tuple_delimiter}"Jordan"{tuple_delimiter}"人物"{tuple_delimiter}"Jordan 对探索有着共同的承诺，并与 Taylor 就设备进行了重要互动。"){record_delimiter}
    ("关系"{tuple_delimiter}"Jordan"{tuple_delimiter}"Cruz"{tuple_delimiter}"Jordan 对探索的承诺是对 Cruz 控制与秩序愿景的反抗。"{tuple_delimiter}"意识形态冲突, 反抗"{tuple_delimiter}5){record_delimiter}
    ("content_keywords"{tuple_delimiter}"power dynamics, ideological conflict, discovery, rebellion"){completion_delimiter}
    ("实体"{tuple_delimiter}"Jordan"{tuple_delimiter}"人物"{tuple_delimiter}"Jordan 对探索有着共同的承诺，并与 Taylor 就设备进行了重要互动。"){record_delimiter}
    glean补充("实体"{tuple_delimiter}"J"{tuple_delimiter}"人物"{tuple_delimiter}"J 对探索有承诺，并与 T互动。"){record_delimiter}
    ("关系"{tuple_delimiter}"Jordan"{tuple_delimiter}"Cruz"{tuple_delimiter}"Jordan 对探索的承诺是对 Cruz 控制与秩序愿景的反抗。"{tuple_delimiter}"意识形态冲突, 反抗"{tuple_delimiter}5){record_delimiter}
    ("content_keywords"{tuple_delimiter}"power dynamics, ideological conflict, discovery, rebellion"){completion_delimiter}

    就会变成

    ["("实体"{tuple_delimiter}"Jordan"...")",
    "("关系"{tuple_delimiter}"Jordan...)",
    ("content_keywords"{tuple_delimiter}"power dynamics,...")",...]
    '''
    records = split_string_by_multi_markers(
        final_result,
        [context_base["record_delimiter"], context_base["completion_delimiter"]],
    )

    # from collections import defaultdict
    """
    maybe_nodes =
    {
        实体1_name->str:
        [dict(
            entity_name=entity_name,
            entity_type=entity_type,
            description=entity_description,
    块id    source_id=entity_source_id, 
        ),...],
        实体2_name->str:
        [dict(
            entity_name=entity_name,
            entity_type=entity_type,
            description=entity_description,
    块id    source_id=entity_source_id, 
        ),...]
    }
    maybe_edges =
    {   
        (实施实体->str, 接收实体->str):
        [dict(
    实施实体         src_id=source,
    接收实体         tgt_id=target,
            weight=weight,
            description=edge_description,
            keywords=edge_keywords,
            source_id=edge_source_id,
        ),...]
    }
    简记为
    maybe_nodes = {str:list[dict]}
    mabybe_edges = {(str,str):list[dict]}
    """
    maybe_nodes = defaultdict(list)
    maybe_edges = defaultdict(list)

    for record in records:
        # 获取括号里内容"("实体"{tuple_delimiter}"Jordan"...")"->""实体"{tuple_delimiter}"Jordan"...""
        record = re.search(r"\((.*)\)", record)
        if record is None:
            continue
        record = record.group(1)
        # record_attributes=List[str,str...]
        record_attributes = split_string_by_multi_markers(
            record, [context_base["tuple_delimiter"]]
        )
        """
        if_entities=
        dict(
            entity_name=entity_name,
            entity_type=entity_type,
            description=entity_description,
    块id    source_id=entity_source_id, 
        )
        """
        if_entities = await _handle_single_entity_extraction(
            record_attributes, chunk_key
        )

        # maybe_nodes为默认 值 为空列表的 字典
        # {JORDAN:[{e_name:..., e_type:.......}]}
        if if_entities is not None:
            maybe_nodes[if_entities["entity_name"]].append(if_entities)
            continue
        """
        if_relation=
        dict(
            src_id=source,
            tgt_id=target,
            weight=weight,
            description=edge_description,
            keywords=edge_keywords,
            source_id=edge_source_id,
        )
        单条record处理出relation
        """
        if_relation = await _handle_single_relationship_extraction(
            record_attributes, chunk_key
        )
        if if_relation is not None:
            # {(JORDAN,JENNY):[{src_id:..., tgt_id:..., }]}
            maybe_edges[(if_relation["src_id"], if_relation["tgt_id"])].append(
                if_relation
            )

    # 下面三个already都是chunk级循环外的
    already_processed += 1
    # 累加每个chunk带来的实体数量和关系数量
    already_entities += len(maybe_nodes)
    already_relations += len(maybe_edges)
    # PROMPTS["process_tickers"] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
    now_ticks = PROMPTS["process_tickers"][
        already_processed % len(PROMPTS["process_tickers"])
    ]
    # 打印处理了多少个chunks，已经处理了的实体和关系
    print(
        f"{now_ticks} Processed {already_processed} chunks, {already_entities} entities(duplicated), {already_relations} relations(duplicated)\r",
        end="",
        flush=True,
    )
    # 由defaultdict转回dict
    return dict(maybe_nodes), dict(maybe_edges)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# =====存实体向量数据=====
if entity_vdb is not None:
    data_for_vdb = {
        compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
            "content": dp["entity_name"] + dp["description"],
            "entity_name": dp["entity_name"],
        }
        for dp in all_entities_data
    }
    # data_for_vdb: dict[str, dict]
    await entity_vdb.upsert(data_for_vdb)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

async def upsert(self, data: dict[str, dict]):
    # 日志和长度处理
    logger.info(f"Inserting {len(data)} vectors to {self.namespace}")
    if not len(data):
        logger.warning("You insert an empty data to vector DB")
        return []

    # k1, v1为dict中的键值对 str:str, 他们被解包出来，最终list_data=List[{str:str}]
    # 实体vdb中 list-data中每个字典元素为
    # {
    # "__id__": hash_id,
    # "content": dp["entity_name"] + dp["description"]
    # "entity_name": dp["entity_name"]
    # }
    list_data = [
        {
            "__id__": k,
            **{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
        }
        for k, v in data.items()
    ]
    # List[str]
    contents = [v["content"] for v in data.values()]

    # List[List[32个str]]
    batches = [
        contents[i : i + self._max_batch_size]
        for i in range(0, len(contents), self._max_batch_size)
    ]
    # 每批32个str输入embedding_func处理成np.array(32, dim)的向量数据
    embedding_tasks = [self.embedding_func(batch) for batch in batches]
    embeddings_list = []
    for f in tqdm_async(
        asyncio.as_completed(embedding_tasks),
        total=len(embedding_tasks),
        desc="Generating embeddings",
        unit="batch",
    ):
        embeddings = await f
        # embeddings_list: List[np.array(32,dim), np.array(32,dim)...]
        # (n,32,dim)
        embeddings_list.append(embeddings)
    # 将多个数组沿着指定的轴（默认是第0轴）连接
    # (n,32,dim)->(n*32,dim) 每行都是一个(实体名+描述)的vector
    embeddings = np.concatenate(embeddings_list)
    # {
    # "__id__": hash_id,
    # "content": dp["entity_name"] + dp["description"]
    # "entity_name": dp["entity_name"]
    # "__vector__": np.array(dim)
    # }
    for i, d in enumerate(list_data):
        d["__vector__"] = embeddings[i]
    # from nano_vectordb import NanoVectorDB
    # NanoVectorDB 的 upsert方法，datas关键字传入List[dict]
    # 将List[dict]中内容写到.json文件中因为之前的创建，以它内置的方式。
    results = self._client.upsert(datas=list_data)
    return results

1
2
3
4
5
6
7
8
9

    async def aquery(self, query: str, param: QueryParam = QueryParam()):
        if param.mode in ["local", "global", "hybrid"]:
            response = await kg_query(...)
        elif param.mode == "naive":
            response = await naive_query(...)
        else:
            raise ValueError(f"Unknown mode {param.mode}")
        await self._query_done()
        return response

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

async def naive_query(
    query,
    chunks_vdb: BaseVectorStorage,
    text_chunks_db: BaseKVStorage[TextChunkSchema],
    query_param: QueryParam,
    global_config: dict,
):
    """
    query: str
    param=QueryParam(mode="naive")
    chunks_vdb: embedding后存入这里，也是字典中包含 chunk的embedding
    text_chunks_db: 原文字符串存入这里，字典结构
    """
    
    use_model_func = global_config["llm_model_func"]
    
    # 一、将提问str转为vector，然后再chunks_vdb中匹配到“cosine_better_than_threshold”的chunk_id
    # 获取排名前top_k个相关文本块
    # top_k: int = 60
    results = await chunks_vdb.query(query, top_k=query_param.top_k)

    if not len(results):
        return PROMPTS["fail_response"]
    
    chunks_ids = [r["id"] for r in results]
    
    # 二、然后根据 chunk_id去 text_chunks_db里找到strs
    # chunks: List[str]
    chunks = await text_chunks_db.get_by_ids(chunks_ids)

    maybe_trun_chunks = truncate_list_by_token_size(
        chunks,
        key=lambda x: x["content"],
        max_token_size=query_param.max_token_for_text_unit,
    )
    logger.info(f"Truncate {len(chunks)} to {len(maybe_trun_chunks)} chunks")
    # section:str = （比如）
    # "机器学习是人工智能的核心...
    # --New Chunk--
    # 深度神经网络通过多层非线性变换...
    # --New Chunk--
    # Transformer架构在NLP领域..."
    section = "\n--New Chunk--\n".join([c["content"] for c in maybe_trun_chunks])
    if query_param.only_need_context:
        return section
    
	# 三、最后将strs嵌入prompt输入LLM
    sys_prompt_temp = PROMPTS["naive_rag_response"]
    sys_prompt = sys_prompt_temp.format(
        content_data=section, response_type=query_param.response_type
    )
    if query_param.only_need_prompt:
        return sys_prompt
    response = await use_model_func(
        query,
        system_prompt=sys_prompt,
    )

    if len(response) > len(sys_prompt):
        response = (
            response[len(sys_prompt) :]
            .replace(sys_prompt, "")
            .replace("user", "")
            .replace("model", "")
            .replace(query, "")
            .replace("<system>", "")
            .replace("</system>", "")
            .strip()
        )

    return response

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

async def kg_query(
    query,
    knowledge_graph_inst: BaseGraphStorage,
    entities_vdb: BaseVectorStorage,
    relationships_vdb: BaseVectorStorage,
    text_chunks_db: BaseKVStorage[TextChunkSchema],
    query_param: QueryParam,
    global_config: dict,
) -> str:
    """
    query: str
    knowledge_graph_inst: NetworkXStorage
    
    """
    
    context = None
    example_number = global_config["addon_params"].get("example_number", None)
    # 1/2/3 个  examples=多个str连接
    if example_number and example_number < len(PROMPTS["keywords_extraction_examples"]):
        examples = "\n".join(
            PROMPTS["keywords_extraction_examples"][: int(example_number)]
        )
    else:
        examples = "\n".join(PROMPTS["keywords_extraction_examples"])
    language = global_config["addon_params"].get(
        "language", PROMPTS["DEFAULT_LANGUAGE"]
    )

    # check mode
    if query_param.mode not in ["local", "global", "hybrid"]:
        logger.error(f"Unknown mode {query_param.mode} in kg_query")
        return PROMPTS["fail_response"]

    # LLM generate keywords
    use_model_func = global_config["llm_model_func"]
    kw_prompt_temp = PROMPTS["keywords_extraction"]
    kw_prompt = kw_prompt_temp.format(query=query, examples=examples, language=language)
    # Output:
    # {{
    #     "high_level_keywords": ["Education", "Poverty reduction", "Socioeconomic development"],
    #     "low_level_keywords": ["School access", "Literacy rates", "Job training", "Income inequality"]
    # }}
    result = await use_model_func(kw_prompt, keyword_extraction=True)
    logger.info("kw_prompt result:")
    print(result)
    try:
        # json_text = locate_json_string_body_from_string(result) # handled in use_model_func
        keywords_data = json.loads(result)
        # List[str]
        hl_keywords = keywords_data.get("high_level_keywords", [])
        ll_keywords = keywords_data.get("low_level_keywords", [])

    # Handle parsing error
    except json.JSONDecodeError as e:
        print(f"JSON parsing error: {e} {result}")
        return PROMPTS["fail_response"]

    # Handdle keywords missing
    # 如果local，ll_keywords必为"kw1, kw2, ..."
    # 如果global，hl_keywords必为"kw1, kw2, ..."
    # 如果hybrid，二者都必为"kw1, kw2, ..."
    if hl_keywords == [] and ll_keywords == []:
        logger.warning("low_level_keywords and high_level_keywords is empty")
        return PROMPTS["fail_response"]
    if ll_keywords == [] and query_param.mode in ["local", "hybrid"]:
        logger.warning("low_level_keywords is empty")
        return PROMPTS["fail_response"]
    else:
        ll_keywords = ", ".join(ll_keywords)
    if hl_keywords == [] and query_param.mode in ["global", "hybrid"]:
        logger.warning("high_level_keywords is empty")
        return PROMPTS["fail_response"]
    else:
        hl_keywords = ", ".join(hl_keywords)

    # Build context
    # ["kw1, kw2, ...", "kw1, kw2, ..."]
    keywords = [ll_keywords, hl_keywords]

    # query_param.mode是不一样的，会选择不同模式，返回的context不一样
    # 貌似需要关键字对应图中结点，万一同义不同词咋办，学校/院校，就对应不上了
    # context:str =
    """
    -----Entities-----
    ```csv
    {entities_context}