Implement SimpleKGBuilder with this setup

neo4j · Dec 3, 2024 · ecda174 · ecda174
1 parent adfc908
commit ecda174
Show file tree

Hide file tree

Showing 4 changed files with 254 additions and 59 deletions.
diff --git a/examples/customize/build_graph/pipeline/simple_kg_pipeline_config.json b/examples/customize/build_graph/pipeline/simple_kg_pipeline_config.json
@@ -1,36 +1,38 @@
 {
     "version_": "1",
+    "template_": "SimpleKGPipeline",
     "neo4j_config": {
-        "uri": {
-            "resolver_": "ENV",
-            "var_": "NEO4J_URI"
-        },
-        "user": {
-            "resolver_": "ENV",
-            "var_": "NEO4J_USER"
-        },
-        "password": {
-            "resolver_": "ENV",
-            "var_": "NEO4J_PASSWORD"
-        },
-        "database": {
-            "resolver_": "ENV",
-            "var_": "NEO4J_DATABASE"
+        "params_": {
+            "uri": {
+                "resolver_": "ENV",
+                "var_": "NEO4J_URI"
+            },
+            "user": {
+                "resolver_": "ENV",
+                "var_": "NEO4J_USER"
+            },
+            "password": {
+                "resolver_": "ENV",
+                "var_": "NEO4J_PASSWORD"
+            }
         }
     },
     "llm_config": {
-        "name_": "openai",
         "class_": "OpenAILLM",
         "params_": {
             "api_key": {
                 "resolver_": "ENV",
                 "var_": "OPENAI_API_KEY"
             },
-            "model_name": "gpt-4o"
+            "model_name": "gpt-4o",
+            "model_params": {
+                "temperature": 0,
+                "max_tokens": 2000,
+                "response_format": {"type": "json_object"}
+            }
         }
     },
     "embedder_config": {
-        "name_": "openai",
         "class_": "OpenAIEmbeddings",
         "params_": {
             "api_key": {
@@ -40,18 +42,71 @@
         }
     },
     "from_pdf": false,
-    "entities": ["Person", {"label":  "Organization"}],
-    "relations": ["WORKS_FOR", {"label":  "DIRECTED_BY"}],
+    "entities": [
+        "Person",
+        {
+            "label": "House",
+            "description": "Family the person belongs to",
+            "properties": [
+                {
+                    "name": "name",
+                    "type": "STRING"
+                }
+            ]
+        },
+        {
+            "label": "Planet",
+            "properties": [
+                {
+                    "name": "name",
+                    "type": "STRING"
+                },
+                {
+                    "name": "weather",
+                    "type": "STRING"
+                }
+            ]
+        }
+    ],
+    "relations": [
+        "PARENT_OF",
+        {
+            "label": "HEIR_OF",
+            "description": "Used for inheritor relationship between father and sons"
+        },
+        {
+            "label": "RULES",
+            "properties": [
+                {
+                    "name": "fromYear",
+                    "type": "INTEGER"
+                }
+            ]
+        }
+    ],
     "potential_schema": [
-        ["Person", "WORKS_FOR", "Organization"],
-        ["Organization", "DIRECTED_BY", "Person"]
+        [
+            "Person",
+            "PARENT_OF",
+            "Person"
+        ],
+        [
+            "Person",
+            "HEIR_OF",
+            "House"
+        ],
+        [
+            "House",
+            "RULES",
+            "Planet"
+        ]
     ],
     "text_splitter": {
-        "class_": "fixed_size_splitter.FixedSizeSplitter",
+        "class_": "text_splitters.fixed_size_splitter.FixedSizeSplitter",
         "params_": {
-            "chunk_size":  100,
+            "chunk_size": 100,
             "chunk_overlap": 10
         }
     },
-    "perform_entity_resolution": false
+    "perform_entity_resolution": true
 }
diff --git a/examples/customize/build_graph/pipeline/simple_kg_pipeline_from_config_file.py b/examples/customize/build_graph/pipeline/simple_kg_pipeline_from_config_file.py
@@ -12,7 +12,7 @@
 # env vars manually set for testing:
 import os
 
-from neo4j_graphrag.experimental.pipeline.config.parser import SimpleKGPipelineBuilder
+from neo4j_graphrag.experimental.pipeline.config.config_poc import PipelineRunner
 from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
 
 os.environ["NEO4J_URI"] = "bolt://localhost:7687"
@@ -28,8 +28,8 @@
 
 async def main() -> PipelineResult:
     file_path = "examples/customize/build_graph/pipeline/simple_kg_pipeline_config.json"
-    pipeline = SimpleKGPipelineBuilder.from_config_file(file_path)
-    return await pipeline.run_async(text=TEXT)
+    pipeline = PipelineRunner.from_config_file(file_path)
+    return await pipeline.run({"text":TEXT})
 
 
 if __name__ == "__main__":

diff --git a/src/neo4j_graphrag/experimental/components/schema.py b/src/neo4j_graphrag/experimental/components/schema.py
@@ -57,8 +57,10 @@ class SchemaEntity(BaseModel):
 
     @classmethod
     def from_text_or_dict(
-        cls, input: str | dict[str, Union[str, dict[str, str]]]
+        cls, input: SchemaEntity | str | dict[str, Union[str, dict[str, str]]]
     ) -> Self:
+        if isinstance(input, SchemaEntity):
+            return input
         if isinstance(input, str):
             return cls(label=input)
         return cls.model_validate(input)
@@ -75,8 +77,10 @@ class SchemaRelation(BaseModel):
 
     @classmethod
     def from_text_or_dict(
-        cls, input: str | dict[str, Union[str, dict[str, str]]]
+        cls, input: SchemaRelation | str | dict[str, Union[str, dict[str, str]]]
     ) -> Self:
+        if isinstance(input, SchemaRelation):
+            return input
         if isinstance(input, str):
             return cls(label=input)
         return cls.model_validate(input)