zrguo commited on
Commit
796af0f
·
2 Parent(s): 538a35a ca2b9e8

Merge pull request #141 from Yazington/feature/other-storage

Browse files
Files changed (2) hide show
  1. lightrag/lightrag.py +3 -1
  2. lightrag/operate.py +11 -11
lightrag/lightrag.py CHANGED
@@ -85,7 +85,9 @@ class LightRAG:
85
 
86
  # LLM
87
  llm_model_func: callable = gpt_4o_mini_complete # hf_model_complete#
88
- llm_model_name: str = "meta-llama/Llama-3.2-1B-Instruct" #'meta-llama/Llama-3.2-1B'#'google/gemma-2-2b-it'
 
 
89
  llm_model_max_token_size: int = 32768
90
  llm_model_max_async: int = 16
91
 
 
85
 
86
  # LLM
87
  llm_model_func: callable = gpt_4o_mini_complete # hf_model_complete#
88
+ llm_model_name: str = (
89
+ "meta-llama/Llama-3.2-1B-Instruct" #'meta-llama/Llama-3.2-1B'#'google/gemma-2-2b-it'
90
+ )
91
  llm_model_max_token_size: int = 32768
92
  llm_model_max_async: int = 16
93
 
lightrag/operate.py CHANGED
@@ -160,7 +160,7 @@ async def _merge_nodes_then_upsert(
160
  description=description,
161
  source_id=source_id,
162
  )
163
- await knwoledge_graph_inst.upsert_node(
164
  entity_name,
165
  node_data=node_data,
166
  )
@@ -172,7 +172,7 @@ async def _merge_edges_then_upsert(
172
  src_id: str,
173
  tgt_id: str,
174
  edges_data: list[dict],
175
- knwoledge_graph_inst: BaseGraphStorage,
176
  global_config: dict,
177
  ):
178
  already_weights = []
@@ -180,8 +180,8 @@ async def _merge_edges_then_upsert(
180
  already_description = []
181
  already_keywords = []
182
 
183
- if await knwoledge_graph_inst.has_edge(src_id, tgt_id):
184
- already_edge = await knwoledge_graph_inst.get_edge(src_id, tgt_id)
185
  already_weights.append(already_edge["weight"])
186
  already_source_ids.extend(
187
  split_string_by_multi_markers(already_edge["source_id"], [GRAPH_FIELD_SEP])
@@ -202,8 +202,8 @@ async def _merge_edges_then_upsert(
202
  set([dp["source_id"] for dp in edges_data] + already_source_ids)
203
  )
204
  for need_insert_id in [src_id, tgt_id]:
205
- if not (await knwoledge_graph_inst.has_node(need_insert_id)):
206
- await knwoledge_graph_inst.upsert_node(
207
  need_insert_id,
208
  node_data={
209
  "source_id": source_id,
@@ -214,7 +214,7 @@ async def _merge_edges_then_upsert(
214
  description = await _handle_entity_relation_summary(
215
  (src_id, tgt_id), description, global_config
216
  )
217
- await knwoledge_graph_inst.upsert_edge(
218
  src_id,
219
  tgt_id,
220
  edge_data=dict(
@@ -237,7 +237,7 @@ async def _merge_edges_then_upsert(
237
 
238
  async def extract_entities(
239
  chunks: dict[str, TextChunkSchema],
240
- knwoledge_graph_inst: BaseGraphStorage,
241
  entity_vdb: BaseVectorStorage,
242
  relationships_vdb: BaseVectorStorage,
243
  global_config: dict,
@@ -341,13 +341,13 @@ async def extract_entities(
341
  maybe_edges[tuple(sorted(k))].extend(v)
342
  all_entities_data = await asyncio.gather(
343
  *[
344
- _merge_nodes_then_upsert(k, v, knwoledge_graph_inst, global_config)
345
  for k, v in maybe_nodes.items()
346
  ]
347
  )
348
  all_relationships_data = await asyncio.gather(
349
  *[
350
- _merge_edges_then_upsert(k[0], k[1], v, knwoledge_graph_inst, global_config)
351
  for k, v in maybe_edges.items()
352
  ]
353
  )
@@ -384,7 +384,7 @@ async def extract_entities(
384
  }
385
  await relationships_vdb.upsert(data_for_vdb)
386
 
387
- return knwoledge_graph_inst
388
 
389
 
390
  async def local_query(
 
160
  description=description,
161
  source_id=source_id,
162
  )
163
+ await knowledge_graph_inst.upsert_node(
164
  entity_name,
165
  node_data=node_data,
166
  )
 
172
  src_id: str,
173
  tgt_id: str,
174
  edges_data: list[dict],
175
+ knowledge_graph_inst: BaseGraphStorage,
176
  global_config: dict,
177
  ):
178
  already_weights = []
 
180
  already_description = []
181
  already_keywords = []
182
 
183
+ if await knowledge_graph_inst.has_edge(src_id, tgt_id):
184
+ already_edge = await knowledge_graph_inst.get_edge(src_id, tgt_id)
185
  already_weights.append(already_edge["weight"])
186
  already_source_ids.extend(
187
  split_string_by_multi_markers(already_edge["source_id"], [GRAPH_FIELD_SEP])
 
202
  set([dp["source_id"] for dp in edges_data] + already_source_ids)
203
  )
204
  for need_insert_id in [src_id, tgt_id]:
205
+ if not (await knowledge_graph_inst.has_node(need_insert_id)):
206
+ await knowledge_graph_inst.upsert_node(
207
  need_insert_id,
208
  node_data={
209
  "source_id": source_id,
 
214
  description = await _handle_entity_relation_summary(
215
  (src_id, tgt_id), description, global_config
216
  )
217
+ await knowledge_graph_inst.upsert_edge(
218
  src_id,
219
  tgt_id,
220
  edge_data=dict(
 
237
 
238
  async def extract_entities(
239
  chunks: dict[str, TextChunkSchema],
240
+ knowledge_graph_inst: BaseGraphStorage,
241
  entity_vdb: BaseVectorStorage,
242
  relationships_vdb: BaseVectorStorage,
243
  global_config: dict,
 
341
  maybe_edges[tuple(sorted(k))].extend(v)
342
  all_entities_data = await asyncio.gather(
343
  *[
344
+ _merge_nodes_then_upsert(k, v, knowledge_graph_inst, global_config)
345
  for k, v in maybe_nodes.items()
346
  ]
347
  )
348
  all_relationships_data = await asyncio.gather(
349
  *[
350
+ _merge_edges_then_upsert(k[0], k[1], v, knowledge_graph_inst, global_config)
351
  for k, v in maybe_edges.items()
352
  ]
353
  )
 
384
  }
385
  await relationships_vdb.upsert(data_for_vdb)
386
 
387
+ return knowledge_graph_inst
388
 
389
 
390
  async def local_query(