Pipeline orchestration with async support

Author: stellasia

poetry.lock

@@ -34,6 +34,7 @@ weaviate-client = {version = "^4.6.1", optional = true}
 pinecone-client = {version = "^4.1.0", optional = true}
 types-mock = "^5.1.0.20240425"
 eval-type-backport = "^0.2.0"
+jsonpath-ng = "^1.6.1"
 
 [tool.poetry.group.dev.dependencies]
 pylint = "^3.1.0"

pyproject.toml

@@ -1,87 +1,74 @@
 import asyncio
+from typing import Any
 from neo4j_genai.core.pipeline import Component
 
 
 class DocumentChunker(Component):
-    def process(self, text: str):
-        return {
-            "chunks": [t.strip() for t in text.split(".") if t.strip()]
-        }
+    async def process(self, text: str) -> dict[str, Any]:
+        return {"chunks": [t.strip() for t in text.split(".") if t.strip()]}
 
 
 class SchemaBuilder(Component):
-    def process(self, schema: dict):
+    async def process(self, schema: dict[str, Any]) -> dict[str, Any]:
         return {"schema": schema}
 
 
 class ERExtractor(Component):
-    def process(self, chunks: list[str], schema: str) -> dict:
+    async def _process_chunk(self, chunk: str, schema: str) -> dict[str, Any]:
         return {
             "data": {
                 "entities": [{"label": "Person", "properties": {"name": "John Doe"}}],
-                "relations": []
+                "relations": [],
             }
         }
 
-    async def _process_chunk(self, chunk: str, schema: str):
-        return {
-            "data": {
-                "entities": [{"label": "Person", "properties": {"name": "John Doe"}}],
-                "relations": []
-            }
-        }
-
-    async def aprocess(self, chunks: list[str], schema: str) -> dict:
-        tasks = [
-            self._process_chunk(chunk, schema)
-            for chunk in chunks
-        ]
+    async def process(self, chunks: list[str], schema: str) -> dict[str, Any]:
+        tasks = [self._process_chunk(chunk, schema) for chunk in chunks]
         result = await asyncio.gather(*tasks)
-        merged_result = {"data": {"entities": [], "relations": []}}
+        merged_result: dict[str, Any] = {"data": {"entities": [], "relations": []}}
         for res in result:
             merged_result["data"]["entities"] += res["data"]["entities"]
             merged_result["data"]["relations"] += res["data"]["relations"]
         return merged_result
 
 
 class Writer(Component):
-    def process(self, entities: dict, relations: dict) -> dict:
+    async def process(
+        self, entities: dict[str, Any], relations: dict[str, Any]
+    ) -> dict[str, Any]:
         return {
             "status": "OK",
             "entities": entities,
             "relations": relations,
         }
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     from neo4j_genai.core.pipeline import Pipeline
 
     pipe = Pipeline()
     pipe.add_component("chunker", DocumentChunker())
-    pipe.add_component("schema",  SchemaBuilder())
+    pipe.add_component("schema", SchemaBuilder())
     pipe.add_component("extractor", ERExtractor())
     pipe.add_component("writer", Writer())
-    pipe.connect("chunker", "extractor", input_defs={
-        "chunks": "chunker.chunks"
-    })
-    pipe.connect("schema", "extractor", input_defs={
-        "schema": "schema.schema"
-    })
-    pipe.connect("extractor", "writer", input_defs={
-        "entities": "extractor.data.entities",
-        "relations": "extractor.data.relations",
-    })
+    pipe.connect("chunker", "extractor", input_defs={"chunks": "chunker.chunks"})
+    pipe.connect("schema", "extractor", input_defs={"schema": "schema.schema"})
+    pipe.connect(
+        "extractor",
+        "writer",
+        input_defs={
+            "entities": "extractor.data.entities",
+            "relations": "extractor.data.relations",
+        },
+    )
 
     pipe_inputs = {
-        "chunker":
-            {
-                "text": "Graphs are everywhere. "
-                        "GraphRAG is the future of Artificial Intelligence. "
-                        "Robots are already running the world."
-            },
-        "schema": {
-            "schema": "Person OWNS House"
-        }
+        "chunker": {
+            "text": "Graphs are everywhere. "
+            "GraphRAG is the future of Artificial Intelligence. "
+            "Robots are already running the world."
+        },
+        "schema": {"schema": "Person OWNS House"},
     }
     # print(pipe.run_all(pipe_inputs))
-    print(asyncio.run(pipe.arun_all(pipe_inputs)))
+    print(asyncio.run(pipe.run(pipe_inputs)))

src/neo4j_genai/components/kg_builder.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any
 
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
@@ -13,18 +14,18 @@ def search(self, *args: Any, **kwargs: Any) -> RetrieverResult:
             ]
         )
 
-    def process(self, query: str):
+    async def process(self, query: str) -> dict[str, Any]:
         res = self.search(query)
         return {"context": "\n".join(c.content for c in res.items)}
 
 
 class PromptTemplate(Component):
-    def process(self, query: str, context: list):
+    async def process(self, query: str, context: list[str]) -> dict[str, Any]:
         return {"prompt": f"my prompt using '{context}', query '{query}'"}
 
 
 class LLM(Component):
-    def process(self, prompt):
+    async def process(self, prompt: str) -> dict[str, Any]:
         return {"answer": f"some text based on '{prompt}'"}
 
 
@@ -37,4 +38,8 @@ def process(self, prompt):
     pipe.connect("augment", "generate", {"prompt": "augment.prompt"})
 
     query = "my question"
-    print(pipe.run_all({"retrieve": {"query": query}, "augment": {"query": query}}))
+    print(
+        asyncio.run(
+            pipe.run({"retrieve": {"query": query}, "augment": {"query": query}})
+        )
+    )

src/neo4j_genai/components/rag.py

@@ -1,40 +1,42 @@
 """
 Basic graph structure for Pipeline
 """
-from typing import Optional
+
+from typing import Optional, Any
+
 
 class Node:
-    def __init__(self, name: str, data: dict) -> None:
+    def __init__(self, name: str, data: dict[str, Any]) -> None:
         self.name = name
         self.data = data
-        self.parents = []
-        self.children = []
+        self.parents: list[Node] = []
+        self.children: list[Node] = []
 
-    def is_root(self):
+    def is_root(self) -> bool:
         return len(self.parents) == 0
 
-    def is_leaf(self):
+    def is_leaf(self) -> bool:
         return len(self.children) == 0
 
 
 class Edge:
-    def __init__(self, start: Node, end: Node, data: Optional[dict] = None):
+    def __init__(self, start: Node, end: Node, data: Optional[dict[str, Any]] = None):
         self.start = start
         self.end = end
         self.data = data
 
 
 class Graph:
-    def __init__(self):
-        self._nodes = {}
-        self._edges = []
+    def __init__(self) -> None:
+        self._nodes: dict[str, Node] = {}
+        self._edges: list[Edge] = []
 
-    def add_node(self, node: Node):
+    def add_node(self, node: Node) -> None:
         if node in self:
             raise ValueError(f"Node {node.name} already exists")
         self._nodes[node.name] = node
 
-    def connect(self, start: Node, end: Node, data: dict):
+    def connect(self, start: Node, end: Node, data: dict[str, Any]) -> None:
         self._edges.append(Edge(start, end, data))
         self._nodes[end.name].parents.append(start)
         self._nodes[start.name].children.append(end)
@@ -43,28 +45,28 @@ def get_node_by_name(self, name: str, raise_exception: bool = False) -> Node:
         node = self._nodes.get(name)
         if node is None and raise_exception:
             raise KeyError(f"Component {name} not in graph")
-        return node
+        return node  # type: ignore
 
-    def roots(self):
+    def roots(self) -> list[Node]:
         root = []
         for node in self._nodes.values():
             if node.is_root():
                 root.append(node)
         return root
 
-    def next_edges(self, node):
+    def next_edges(self, node: Node) -> list[Edge]:
         res = []
         for edge in self._edges:
             if edge.start == node:
                 res.append(edge)
         return res
 
-    def previous_edges(self, node):
+    def previous_edges(self, node: Node) -> list[Edge]:
         res = []
         for edge in self._edges:
             if edge.end == node:
                 res.append(edge)
         return res
 
-    def __contains__(self, node: Node):
+    def __contains__(self, node: Node) -> bool:
         return node.name in self._nodes