Remove SchemaEntity/SchemaRelation as valid inputs in SimpleKGPipeline

(not "simple")
neo4j · Nov 5, 2024 · 71b0eda · 71b0eda
1 parent f02fa10
commit 71b0eda
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/neo4j_graphrag/experimental/pipeline/kg_builder.py b/src/neo4j_graphrag/experimental/pipeline/kg_builder.py
@@ -65,8 +65,8 @@ class SimpleKGPipelineConfig(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
-EntityInputType = Union[str, dict[str, str], SchemaEntity]
-RelationInputType = Union[str, dict[str, str], SchemaRelation]
+EntityInputType = Union[str, dict[str, Any]]
+RelationInputType = Union[str, dict[str, Any]]
 
 
 class SimpleKGPipeline:
@@ -166,16 +166,12 @@ def __init__(
 
     @staticmethod
     def to_schema_entity(entity: EntityInputType) -> SchemaEntity:
-        if isinstance(entity, SchemaEntity):
-            return entity
         if isinstance(entity, dict):
             return SchemaEntity(**entity)
         return SchemaEntity(label=entity)
 
     @staticmethod
     def to_schema_relation(relation: RelationInputType) -> SchemaRelation:
-        if isinstance(relation, SchemaRelation):
-            return relation
         if isinstance(relation, dict):
             return SchemaRelation(**relation)
         return SchemaRelation(label=relation)