Merge pull request #141 from Yazington/feature/other-storage
Browse files- lightrag/lightrag.py +3 -1
- 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 =
|
|
|
|
|
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
|
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 |
-
|
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
|
184 |
-
already_edge = await
|
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
|
206 |
-
await
|
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
|
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 |
-
|
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,
|
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,
|
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
|
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(
|