code clean
Browse files- lightrag/lightrag.py +18 -21
lightrag/lightrag.py
CHANGED
@@ -536,7 +536,7 @@ class LightRAG:
|
|
536 |
),
|
537 |
embedding_func=self.embedding_func,
|
538 |
)
|
539 |
-
|
540 |
self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
|
541 |
partial(
|
542 |
self.llm_model_func, # type: ignore
|
@@ -955,37 +955,34 @@ class LightRAG:
|
|
955 |
|
956 |
# Insert entities into vector storage if needed
|
957 |
data_for_vdb = {
|
958 |
-
|
959 |
-
|
960 |
-
|
961 |
-
}
|
962 |
-
for dp in all_entities_data
|
963 |
}
|
|
|
|
|
964 |
await self.entities_vdb.upsert(data_for_vdb)
|
965 |
|
966 |
# Insert relationships into vector storage if needed
|
967 |
data_for_vdb = {
|
968 |
-
|
969 |
-
|
970 |
-
|
971 |
-
|
972 |
-
|
973 |
-
|
974 |
-
|
975 |
-
}
|
976 |
-
for dp in all_relationships_data
|
977 |
}
|
|
|
|
|
978 |
await self.relationships_vdb.upsert(data_for_vdb)
|
979 |
-
|
980 |
finally:
|
981 |
if update_storage:
|
982 |
await self._insert_done()
|
983 |
|
984 |
def query(
|
985 |
-
self,
|
986 |
-
query: str,
|
987 |
-
param: QueryParam = QueryParam(),
|
988 |
-
prompt: str | None = None
|
989 |
) -> str:
|
990 |
"""
|
991 |
Perform a sync query.
|
@@ -997,7 +994,7 @@ class LightRAG:
|
|
997 |
|
998 |
Returns:
|
999 |
str: The result of the query execution.
|
1000 |
-
"""
|
1001 |
loop = always_get_an_event_loop()
|
1002 |
return loop.run_until_complete(self.aquery(query, param, prompt))
|
1003 |
|
|
|
536 |
),
|
537 |
embedding_func=self.embedding_func,
|
538 |
)
|
539 |
+
|
540 |
self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
|
541 |
partial(
|
542 |
self.llm_model_func, # type: ignore
|
|
|
955 |
|
956 |
# Insert entities into vector storage if needed
|
957 |
data_for_vdb = {
|
958 |
+
compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
|
959 |
+
"content": dp["entity_name"] + dp["description"],
|
960 |
+
"entity_name": dp["entity_name"],
|
|
|
|
|
961 |
}
|
962 |
+
for dp in all_entities_data
|
963 |
+
}
|
964 |
await self.entities_vdb.upsert(data_for_vdb)
|
965 |
|
966 |
# Insert relationships into vector storage if needed
|
967 |
data_for_vdb = {
|
968 |
+
compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): {
|
969 |
+
"src_id": dp["src_id"],
|
970 |
+
"tgt_id": dp["tgt_id"],
|
971 |
+
"content": dp["keywords"]
|
972 |
+
+ dp["src_id"]
|
973 |
+
+ dp["tgt_id"]
|
974 |
+
+ dp["description"],
|
|
|
|
|
975 |
}
|
976 |
+
for dp in all_relationships_data
|
977 |
+
}
|
978 |
await self.relationships_vdb.upsert(data_for_vdb)
|
979 |
+
|
980 |
finally:
|
981 |
if update_storage:
|
982 |
await self._insert_done()
|
983 |
|
984 |
def query(
|
985 |
+
self, query: str, param: QueryParam = QueryParam(), prompt: str | None = None
|
|
|
|
|
|
|
986 |
) -> str:
|
987 |
"""
|
988 |
Perform a sync query.
|
|
|
994 |
|
995 |
Returns:
|
996 |
str: The result of the query execution.
|
997 |
+
"""
|
998 |
loop = always_get_an_event_loop()
|
999 |
return loop.run_until_complete(self.aquery(query, param, prompt))
|
1000 |
|