cleanup

py/cli/__init__.py

@@ -1,6 +1,5 @@
 from .cli import cli, main
 from .command_group import cli as command_group_cli
-
 from .commands import auth, ingestion, management, retrieval, server
 
 __all__ = [

py/core/main/hatchet/restructure_workflow.py

@@ -169,7 +169,9 @@ def __init__(self, restructure_service: RestructureService):
     @r2r_hatchet.step(retries=3, timeout="60m")
     async def kg_node_creation(self, context: Context) -> None:
         input_data = context.workflow_input()["request"]
-        max_description_input_length = input_data["max_description_input_length"]
+        max_description_input_length = input_data[
+            "max_description_input_length"
+        ]
         await self.restructure_service.kg_node_creation(
             max_description_input_length=max_description_input_length
         )

py/core/pipes/kg/node_extraction.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+import random
 from typing import Any, AsyncGenerator, Optional
 from uuid import UUID
 
@@ -17,8 +18,6 @@
 from core.base.abstractions.graph import Entity, Triple
 from core.base.pipes.base_pipe import AsyncPipe
 
-import random
-
 logger = logging.getLogger(__name__)
 
 
@@ -127,7 +126,9 @@ async def _run_logic(
             Ensure the summary is coherent, informative, and captures the essence of the entity within the context of the provided information.
         """
 
-        async def process_entity(entity, triples, max_description_input_length):
+        async def process_entity(
+            entity, triples, max_description_input_length
+        ):
 
             # if embedding is present in the entity, just return it
             # in the future disable this to override and recompute the descriptions for all entities
@@ -136,8 +137,8 @@ async def process_entity(entity, triples, max_description_input_length):
 
             entity_info = f"{entity.name}, {entity.description}"
             triples_txt = [
-                    f"{i+1}: {triple.subject}, {triple.object}, {triple.predicate} - Summary: {triple.description}"
-                    for i, triple in enumerate(triples)
+                f"{i+1}: {triple.subject}, {triple.object}, {triple.predicate} - Summary: {triple.description}"
+                for i, triple in enumerate(triples)
             ]
 
             # truncate the descriptions to the max_description_input_length
@@ -196,13 +197,21 @@ async def process_entity(entity, triples, max_description_input_length):
 
             return out_entity
 
-        max_description_input_length = input.message["max_description_input_length"]
+        max_description_input_length = input.message[
+            "max_description_input_length"
+        ]
         node_extrations = input.message["node_extrations"]
 
         tasks = []
         count = 0
         async for entity, triples in node_extrations:
-            tasks.append(asyncio.create_task(process_entity(entity, triples, max_description_input_length)))
+            tasks.append(
+                asyncio.create_task(
+                    process_entity(
+                        entity, triples, max_description_input_length
+                    )
+                )
+            )
             count += 1
 
         logger.info(f"KG Node Description pipe: Created {count} tasks")