YanSte commited on
Commit
e336d3f
·
1 Parent(s): 6c5270a

cleaned typing

Browse files
Files changed (1) hide show
  1. lightrag/lightrag.py +17 -17
lightrag/lightrag.py CHANGED
@@ -568,7 +568,7 @@ class LightRAG:
568
  input: str | list[str],
569
  split_by_character: str | None = None,
570
  split_by_character_only: bool = False,
571
- ):
572
  """Sync Insert documents with checkpoint support
573
 
574
  Args:
@@ -578,7 +578,7 @@ class LightRAG:
578
  split_by_character is None, this parameter is ignored.
579
  """
580
  loop = always_get_an_event_loop()
581
- return loop.run_until_complete(
582
  self.ainsert(input, split_by_character, split_by_character_only)
583
  )
584
 
@@ -587,7 +587,7 @@ class LightRAG:
587
  input: str | list[str],
588
  split_by_character: str | None = None,
589
  split_by_character_only: bool = False,
590
- ):
591
  """Async Insert documents with checkpoint support
592
 
593
  Args:
@@ -601,13 +601,13 @@ class LightRAG:
601
  split_by_character, split_by_character_only
602
  )
603
 
604
- def insert_custom_chunks(self, full_text: str, text_chunks: list[str]):
605
  loop = always_get_an_event_loop()
606
- return loop.run_until_complete(
607
- self.ainsert_custom_chunks(full_text, text_chunks)
608
- )
609
 
610
- async def ainsert_custom_chunks(self, full_text: str, text_chunks: list[str]):
 
 
611
  update_storage = False
612
  try:
613
  doc_key = compute_mdhash_id(full_text.strip(), prefix="doc-")
@@ -653,7 +653,7 @@ class LightRAG:
653
  if update_storage:
654
  await self._insert_done()
655
 
656
- async def apipeline_enqueue_documents(self, input: str | list[str]):
657
  """
658
  Pipeline for Processing Documents
659
 
@@ -832,7 +832,7 @@ class LightRAG:
832
  logger.error("Failed to extract entities and relationships")
833
  raise e
834
 
835
- async def _insert_done(self):
836
  tasks = [
837
  cast(StorageNameSpace, storage_inst).index_done_callback()
838
  for storage_inst in [ # type: ignore
@@ -848,11 +848,11 @@ class LightRAG:
848
  ]
849
  await asyncio.gather(*tasks)
850
 
851
- def insert_custom_kg(self, custom_kg: dict[str, Any]):
852
  loop = always_get_an_event_loop()
853
- return loop.run_until_complete(self.ainsert_custom_kg(custom_kg))
854
 
855
- async def ainsert_custom_kg(self, custom_kg: dict[str, Any]):
856
  update_storage = False
857
  try:
858
  # Insert chunks into vector storage
@@ -1205,11 +1205,11 @@ class LightRAG:
1205
  async def _query_done(self):
1206
  await self.llm_response_cache.index_done_callback()
1207
 
1208
- def delete_by_entity(self, entity_name: str):
1209
  loop = always_get_an_event_loop()
1210
  return loop.run_until_complete(self.adelete_by_entity(entity_name))
1211
 
1212
- async def adelete_by_entity(self, entity_name: str):
1213
  entity_name = f'"{entity_name.upper()}"'
1214
 
1215
  try:
@@ -1224,7 +1224,7 @@ class LightRAG:
1224
  except Exception as e:
1225
  logger.error(f"Error while deleting entity '{entity_name}': {e}")
1226
 
1227
- async def _delete_by_entity_done(self):
1228
  await asyncio.gather(
1229
  *[
1230
  cast(StorageNameSpace, storage_inst).index_done_callback()
@@ -1497,7 +1497,7 @@ class LightRAG:
1497
 
1498
  async def get_relation_info(
1499
  self, src_entity: str, tgt_entity: str, include_vector_data: bool = False
1500
- ):
1501
  """Get detailed information of a relationship
1502
 
1503
  Args:
 
568
  input: str | list[str],
569
  split_by_character: str | None = None,
570
  split_by_character_only: bool = False,
571
+ ) -> None:
572
  """Sync Insert documents with checkpoint support
573
 
574
  Args:
 
578
  split_by_character is None, this parameter is ignored.
579
  """
580
  loop = always_get_an_event_loop()
581
+ loop.run_until_complete(
582
  self.ainsert(input, split_by_character, split_by_character_only)
583
  )
584
 
 
587
  input: str | list[str],
588
  split_by_character: str | None = None,
589
  split_by_character_only: bool = False,
590
+ ) -> None:
591
  """Async Insert documents with checkpoint support
592
 
593
  Args:
 
601
  split_by_character, split_by_character_only
602
  )
603
 
604
+ def insert_custom_chunks(self, full_text: str, text_chunks: list[str]) -> None:
605
  loop = always_get_an_event_loop()
606
+ loop.run_until_complete(self.ainsert_custom_chunks(full_text, text_chunks))
 
 
607
 
608
+ async def ainsert_custom_chunks(
609
+ self, full_text: str, text_chunks: list[str]
610
+ ) -> None:
611
  update_storage = False
612
  try:
613
  doc_key = compute_mdhash_id(full_text.strip(), prefix="doc-")
 
653
  if update_storage:
654
  await self._insert_done()
655
 
656
+ async def apipeline_enqueue_documents(self, input: str | list[str]) -> None:
657
  """
658
  Pipeline for Processing Documents
659
 
 
832
  logger.error("Failed to extract entities and relationships")
833
  raise e
834
 
835
+ async def _insert_done(self) -> None:
836
  tasks = [
837
  cast(StorageNameSpace, storage_inst).index_done_callback()
838
  for storage_inst in [ # type: ignore
 
848
  ]
849
  await asyncio.gather(*tasks)
850
 
851
+ def insert_custom_kg(self, custom_kg: dict[str, Any]) -> None:
852
  loop = always_get_an_event_loop()
853
+ loop.run_until_complete(self.ainsert_custom_kg(custom_kg))
854
 
855
+ async def ainsert_custom_kg(self, custom_kg: dict[str, Any]) -> None:
856
  update_storage = False
857
  try:
858
  # Insert chunks into vector storage
 
1205
  async def _query_done(self):
1206
  await self.llm_response_cache.index_done_callback()
1207
 
1208
+ def delete_by_entity(self, entity_name: str) -> None:
1209
  loop = always_get_an_event_loop()
1210
  return loop.run_until_complete(self.adelete_by_entity(entity_name))
1211
 
1212
+ async def adelete_by_entity(self, entity_name: str) -> None:
1213
  entity_name = f'"{entity_name.upper()}"'
1214
 
1215
  try:
 
1224
  except Exception as e:
1225
  logger.error(f"Error while deleting entity '{entity_name}': {e}")
1226
 
1227
+ async def _delete_by_entity_done(self) -> None:
1228
  await asyncio.gather(
1229
  *[
1230
  cast(StorageNameSpace, storage_inst).index_done_callback()
 
1497
 
1498
  async def get_relation_info(
1499
  self, src_entity: str, tgt_entity: str, include_vector_data: bool = False
1500
+ ) -> dict[str, str | None | dict[str, str]]:
1501
  """Get detailed information of a relationship
1502
 
1503
  Args: