KG Creation Pipeline Speed Up (#117)

* Made LLM calls concurrent in KG creation pipeline

* Renamed test_kb_builder_pipeline_e2e.py to test_kg_builder_pipeline_e2e.py

* Added an asynchronous Neo4jWriter

* Merged async and non-async Neo4j writers

* Fixed neo4j writer docstring

* Refactored Neo4jWriter

Author: alexthomas93

docs/source/api.rst

@@ -213,6 +213,9 @@ Database Interaction
 
 .. autofunction:: neo4j_genai.indexes.upsert_vector_on_relationship
 
+.. autofunction:: neo4j_genai.indexes.async_upsert_vector
+
+.. autofunction:: neo4j_genai.indexes.async_upsert_vector_on_relationship
 
 ******
 Errors

src/neo4j_genai/experimental/components/entity_relation_extractor.py

@@ -232,6 +232,7 @@ class LLMEntityRelationExtractor(EntityRelationExtractor):
         prompt_template (ERExtractionTemplate | str): A custom prompt template to use for extraction.
         create_lexical_graph (bool): Whether to include the text chunks in the graph in addition to the extracted entities and relations. Defaults to True.
         on_error (OnError): What to do when an error occurs during extraction. Defaults to raising an error.
+        max_concurrency (int): The maximum number of concurrent tasks which can be used to make requests to the LLM.
 
     Example:
 
@@ -255,9 +256,11 @@ def __init__(
         prompt_template: ERExtractionTemplate | str = ERExtractionTemplate(),
         create_lexical_graph: bool = True,
         on_error: OnError = OnError.RAISE,
+        max_concurrency: int = 5,
     ) -> None:
         super().__init__(on_error=on_error, create_lexical_graph=create_lexical_graph)
         self.llm = llm  # with response_format={ "type": "json_object" },
+        self.max_concurrency = max_concurrency
         if isinstance(prompt_template, str):
             template = PromptTemplate(prompt_template, expected_inputs=[])
         else:
@@ -271,7 +274,7 @@ async def extract_for_chunk(
         prompt = self.prompt_template.format(
             text=chunk.text, schema=schema.model_dump(), examples=examples
         )
-        llm_result = self.llm.invoke(prompt)
+        llm_result = await self.llm.ainvoke(prompt)
         try:
             result = json.loads(llm_result.content)
         except json.JSONDecodeError:
@@ -322,12 +325,20 @@ def combine_chunk_graphs(self, chunk_graphs: List[Neo4jGraph]) -> Neo4jGraph:
         return graph
 
     async def run_for_chunk(
-        self, schema: SchemaConfig, examples: str, chunk_index: int, chunk: TextChunk
+        self,
+        schema: SchemaConfig,
+        examples: str,
+        chunk_index: int,
+        chunk: TextChunk,
+        sem: asyncio.Semaphore,
     ) -> Neo4jGraph:
         """Run extraction and post processing for a single chunk"""
-        chunk_graph = await self.extract_for_chunk(schema, examples, chunk_index, chunk)
-        await self.post_process_chunk(chunk_graph, chunk_index, chunk)
-        return chunk_graph
+        async with sem:
+            chunk_graph = await self.extract_for_chunk(
+                schema, examples, chunk_index, chunk
+            )
+            await self.post_process_chunk(chunk_graph, chunk_index, chunk)
+            return chunk_graph
 
     @validate_call
     async def run(
@@ -341,8 +352,9 @@ async def run(
         schema = schema or SchemaConfig(entities={}, relations={}, potential_schema=[])
         examples = examples or ""
         self._id_prefix = str(datetime.now().timestamp())
+        sem = asyncio.Semaphore(self.max_concurrency)
         tasks = [
-            self.run_for_chunk(schema, examples, chunk_index, chunk)
+            self.run_for_chunk(schema, examples, chunk_index, chunk, sem)
             for chunk_index, chunk in enumerate(chunks.chunks)
         ]
         chunk_graphs = await asyncio.gather(*tasks)

src/neo4j_genai/experimental/components/kg_writer.py

@@ -14,9 +14,10 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import asyncio
 import logging
 from abc import abstractmethod
-from typing import Literal, Optional
+from typing import Any, Dict, Literal, Optional, Tuple
 
 import neo4j
 from pydantic import validate_call
@@ -27,7 +28,12 @@
     Neo4jRelationship,
 )
 from neo4j_genai.experimental.pipeline.component import Component, DataModel
-from neo4j_genai.indexes import upsert_vector, upsert_vector_on_relationship
+from neo4j_genai.indexes import (
+    async_upsert_vector,
+    async_upsert_vector_on_relationship,
+    upsert_vector,
+    upsert_vector_on_relationship,
+)
 from neo4j_genai.neo4j_queries import UPSERT_NODE_QUERY, UPSERT_RELATIONSHIP_QUERY
 
 logger = logging.getLogger(__name__)
@@ -64,20 +70,21 @@ class Neo4jWriter(KGWriter):
     Args:
         driver (neo4j.driver): The Neo4j driver to connect to the database.
         neo4j_database (Optional[str]): The name of the Neo4j database to write to. Defaults to 'neo4j' if not provided.
+        max_concurrency (int): The maximum number of concurrent tasks which can be used to make requests to the LLM.
 
     Example:
 
     .. code-block:: python
 
-        from neo4j import GraphDatabase
+        from neo4j import AsyncGraphDatabase
         from neo4j_genai.experimental.components.kg_writer import Neo4jWriter
         from neo4j_genai.experimental.pipeline import Pipeline
 
         URI = "neo4j://localhost:7687"
         AUTH = ("neo4j", "password")
         DATABASE = "neo4j"
 
-        driver = GraphDatabase.driver(URI, auth=AUTH, database=DATABASE)
+        driver = AsyncGraphDatabase.driver(URI, auth=AUTH, database=DATABASE)
         writer = Neo4jWriter(driver=driver, neo4j_database=DATABASE)
 
         pipeline = Pipeline()
@@ -89,16 +96,13 @@ def __init__(
         self,
         driver: neo4j.driver,
         neo4j_database: Optional[str] = None,
+        max_concurrency: int = 5,
     ):
         self.driver = driver
         self.neo4j_database = neo4j_database
+        self.max_concurrency = max_concurrency
 
-    def _upsert_node(self, node: Neo4jNode) -> None:
-        """Upserts a single node into the Neo4j database."
-
-        Args:
-            node (Neo4jNode): The node to upsert into the database.
-        """
+    def _get_node_query(self, node: Neo4jNode) -> Tuple[str, Dict[str, Any]]:
         # Create the initial node
         parameters = {"id": node.id}
         if node.properties:
@@ -107,6 +111,15 @@ def _upsert_node(self, node: Neo4jNode) -> None:
             "{" + ", ".join(f"{key}: ${key}" for key in parameters.keys()) + "}"
         )
         query = UPSERT_NODE_QUERY.format(label=node.label, properties=properties)
+        return query, parameters
+
+    def _upsert_node(self, node: Neo4jNode) -> None:
+        """Upserts a single node into the Neo4j database."
+
+        Args:
+            node (Neo4jNode): The node to upsert into the database.
+        """
+        query, parameters = self._get_node_query(node)
         result = self.driver.execute_query(query, parameters_=parameters)
         node_id = result.records[0]["elementID(n)"]
         # Add the embedding properties to the node
@@ -120,12 +133,32 @@ def _upsert_node(self, node: Neo4jNode) -> None:
                     neo4j_database=self.neo4j_database,
                 )
 
-    def _upsert_relationship(self, rel: Neo4jRelationship) -> None:
-        """Upserts a single relationship into the Neo4j database.
+    async def _async_upsert_node(
+        self,
+        node: Neo4jNode,
+        sem: asyncio.Semaphore,
+    ) -> None:
+        """Asynchronously upserts a single node into the Neo4j database."
 
         Args:
-            rel (Neo4jRelationship): The relationship to upsert into the database.
+            node (Neo4jNode): The node to upsert into the database.
         """
+        async with sem:
+            query, parameters = self._get_node_query(node)
+            result = await self.driver.execute_query(query, parameters_=parameters)
+            node_id = result.records[0]["elementID(n)"]
+            # Add the embedding properties to the node
+            if node.embedding_properties:
+                for prop, vector in node.embedding_properties.items():
+                    await async_upsert_vector(
+                        driver=self.driver,
+                        node_id=node_id,
+                        embedding_property=prop,
+                        vector=vector,
+                        neo4j_database=self.neo4j_database,
+                    )
+
+    def _get_rel_query(self, rel: Neo4jRelationship) -> Tuple[str, Dict[str, Any]]:
         # Create the initial relationship
         parameters = {
             "start_node_id": rel.start_node_id,
@@ -142,6 +175,15 @@ def _upsert_relationship(self, rel: Neo4jRelationship) -> None:
             type=rel.type,
             properties=properties,
         )
+        return query, parameters
+
+    def _upsert_relationship(self, rel: Neo4jRelationship) -> None:
+        """Upserts a single relationship into the Neo4j database.
+
+        Args:
+            rel (Neo4jRelationship): The relationship to upsert into the database.
+        """
+        query, parameters = self._get_rel_query(rel)
         result = self.driver.execute_query(query, parameters_=parameters)
         rel_id = result.records[0]["elementID(r)"]
         # Add the embedding properties to the relationship
@@ -155,6 +197,29 @@ def _upsert_relationship(self, rel: Neo4jRelationship) -> None:
                     neo4j_database=self.neo4j_database,
                 )
 
+    async def _async_upsert_relationship(
+        self, rel: Neo4jRelationship, sem: asyncio.Semaphore
+    ) -> None:
+        """Asynchronously upserts a single relationship into the Neo4j database.
+
+        Args:
+            rel (Neo4jRelationship): The relationship to upsert into the database.
+        """
+        async with sem:
+            query, parameters = self._get_rel_query(rel)
+            result = await self.driver.execute_query(query, parameters_=parameters)
+            rel_id = result.records[0]["elementID(r)"]
+            # Add the embedding properties to the relationship
+            if rel.embedding_properties:
+                for prop, vector in rel.embedding_properties.items():
+                    await async_upsert_vector_on_relationship(
+                        driver=self.driver,
+                        rel_id=rel_id,
+                        embedding_property=prop,
+                        vector=vector,
+                        neo4j_database=self.neo4j_database,
+                    )
+
     @validate_call
     async def run(self, graph: Neo4jGraph) -> KGWriterModel:
         """Upserts a knowledge graph into a Neo4j database.
@@ -163,11 +228,24 @@ async def run(self, graph: Neo4jGraph) -> KGWriterModel:
             graph (Neo4jGraph): The knowledge graph to upsert into the database.
         """
         try:
-            for node in graph.nodes:
-                self._upsert_node(node)
-
-            for rel in graph.relationships:
-                self._upsert_relationship(rel)
+            if isinstance(self.driver, neo4j.AsyncDriver):
+                sem = asyncio.Semaphore(self.max_concurrency)
+                node_tasks = [
+                    self._async_upsert_node(node, sem) for node in graph.nodes
+                ]
+                await asyncio.gather(*node_tasks)
+
+                rel_tasks = [
+                    self._async_upsert_relationship(rel, sem)
+                    for rel in graph.relationships
+                ]
+                await asyncio.gather(*rel_tasks)
+            else:
+                for node in graph.nodes:
+                    self._upsert_node(node)
+
+                for rel in graph.relationships:
+                    self._upsert_relationship(rel)
 
             return KGWriterModel(status="SUCCESS")
         except neo4j.exceptions.ClientError as e: