formatting

py/core/pipes/kg/extraction.py

@@ -190,31 +190,34 @@ async def _run_logic(
     ) -> AsyncGenerator[Union[KGExtraction, R2RDocumentProcessingError], None]:
 
         logger.info("Running KG Extraction Pipe")
-    
+
         async def process_extraction(extraction):
             return await self.extract_kg(extraction)
-    
+
         document_ids = []
         async for extraction in input.message:
             document_ids.append(extraction)
 
         if document_ids == []:
-            document_ids = [doc.id for doc in self.database_provider.relational.get_documents_overview()]
+            document_ids = [
+                doc.id
+                for doc in self.database_provider.relational.get_documents_overview()
+            ]
 
         for document_id in document_ids:
             logger.info(f"Extracting KG for document: {document_id}")
             extractions = [
                 DocumentFragment(
-                    id=extraction['fragment_id'],
-                    extraction_id=extraction['extraction_id'],
-                    document_id=extraction['document_id'],
-                    user_id=extraction['user_id'],
-                    group_ids=extraction['group_ids'],
-                    data=extraction['text'],
-                    metadata=extraction['metadata']
+                    id=extraction["fragment_id"],
+                    extraction_id=extraction["extraction_id"],
+                    document_id=extraction["document_id"],
+                    user_id=extraction["user_id"],
+                    group_ids=extraction["group_ids"],
+                    data=extraction["text"],
+                    metadata=extraction["metadata"],
                 )
                 for extraction in self.database_provider.vector.get_document_chunks(
-                    document_id=document_id 
+                    document_id=document_id
                 )
             ]
 

py/core/pipes/kg/node_extraction.py

@@ -132,7 +132,6 @@ 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):
 
             # if embedding is present in the entity, just return it

py/core/pipes/kg/storage.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Any, AsyncGenerator, Optional
+from typing import Any, AsyncGenerator, Optional, Tuple
 from uuid import UUID
 
 from core.base import (
@@ -12,7 +12,7 @@
 )
 from core.base.pipes.base_pipe import AsyncPipe
 from core.base.providers import KGProvider
-from typing import Tuple
+
 logger = logging.getLogger(__name__)
 
 
@@ -57,7 +57,9 @@ async def store(
         Stores a batch of knowledge graph extractions in the graph database.
         """
         try:
-            return self.kg_provider.upsert_nodes_and_relationships(kg_extractions)
+            return self.kg_provider.upsert_nodes_and_relationships(
+                kg_extractions
+            )
         except Exception as e:
             error_message = f"Failed to store knowledge graph extractions in the database: {e}"
             logger.error(error_message)