Update SimpleKGPipeline for automatic schema extraction

Author: NathalieCharbel

src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py

@@ -32,6 +32,7 @@
     SchemaBuilder,
     SchemaEntity,
     SchemaRelation,
+    SchemaFromText,
 )
 from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter
 from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
@@ -80,13 +81,18 @@ class SimpleKGPipelineConfig(TemplatePipelineConfig):
     perform_entity_resolution: bool = True
     lexical_graph_config: Optional[LexicalGraphConfig] = None
     neo4j_database: Optional[str] = None
+    auto_schema_extraction: bool = False
 
     pdf_loader: Optional[ComponentType] = None
     kg_writer: Optional[ComponentType] = None
     text_splitter: Optional[ComponentType] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    def has_user_provided_schema(self) -> bool:
+        """Check if the user has provided schema information"""
+        return bool(self.entities or self.relations or self.potential_schema)
+
     def _get_pdf_loader(self) -> Optional[PdfLoader]:
         if not self.from_pdf:
             return None
@@ -114,15 +120,26 @@ def _get_run_params_for_splitter(self) -> dict[str, Any]:
     def _get_chunk_embedder(self) -> TextChunkEmbedder:
         return TextChunkEmbedder(embedder=self.get_default_embedder())
 
-    def _get_schema(self) -> SchemaBuilder:
+    def _get_schema(self) -> Union[SchemaBuilder, SchemaFromText]:
+        """
+        Get the appropriate schema component based on configuration.
+        Return SchemaFromText for automatic extraction or SchemaBuilder for manual schema.
+        """
+        if self.auto_schema_extraction and not self.has_user_provided_schema():
+            return SchemaFromText(llm=self.get_default_llm())
         return SchemaBuilder()
 
     def _get_run_params_for_schema(self) -> dict[str, Any]:
-        return {
-            "entities": [SchemaEntity.from_text_or_dict(e) for e in self.entities],
-            "relations": [SchemaRelation.from_text_or_dict(r) for r in self.relations],
-            "potential_schema": self.potential_schema,
-        }
+        if self.auto_schema_extraction and not self.has_user_provided_schema():
+            # for automatic extraction, the text parameter is needed (will flow through the pipeline connections)
+            return {}
+        else:
+            # for manual schema, use the provided entities/relations/potential_schema
+            return {
+                "entities": [SchemaEntity.from_text_or_dict(e) for e in self.entities],
+                "relations": [SchemaRelation.from_text_or_dict(r) for r in self.relations],
+                "potential_schema": self.potential_schema,
+            }
 
     def _get_extractor(self) -> EntityRelationExtractor:
         return LLMEntityRelationExtractor(
@@ -163,6 +180,17 @@ def _get_connections(self) -> list[ConnectionDefinition]:
                     input_config={"text": "pdf_loader.text"},
                 )
             )
+            
+            # handle automatic schema extraction
+            if self.auto_schema_extraction and not self.has_user_provided_schema():
+                connections.append(
+                    ConnectionDefinition(
+                        start="pdf_loader",
+                        end="schema",
+                        input_config={"text": "pdf_loader.text"},
+                    )
+                )
+            
             connections.append(
                 ConnectionDefinition(
                     start="schema",
@@ -174,13 +202,21 @@ def _get_connections(self) -> list[ConnectionDefinition]:
                 )
             )
         else:
+            # handle automatic schema extraction for direct text input: ensure schema extraction uses the complete text
+            if self.auto_schema_extraction and not self.has_user_provided_schema():
+                connections.append(
+                    ConnectionDefinition(
+                        start="__input__",  # connection to pipeline input
+                        end="schema",
+                        input_config={"text": "text"},  # use the original text input
+                    )
+                )
+            
             connections.append(
                 ConnectionDefinition(
                     start="schema",
                     end="extractor",
-                    input_config={
-                        "schema": "schema",
-                    },
+                    input_config={"schema": "schema"},
                 )
             )
         connections.append(

src/neo4j_graphrag/experimental/pipeline/kg_builder.py

@@ -16,6 +16,8 @@
 from __future__ import annotations
 
 from typing import List, Optional, Sequence, Union
+import logging
+import warnings
 
 import neo4j
 from pydantic import ValidationError
@@ -42,7 +44,9 @@
 )
 from neo4j_graphrag.generation.prompts import ERExtractionTemplate
 from neo4j_graphrag.llm.base import LLMInterface
+from neo4j_graphrag.experimental.components.schema import SchemaConfig, SchemaBuilder
 
+logger = logging.getLogger(__name__)
 
 class SimpleKGPipeline:
     """
@@ -53,17 +57,20 @@ class SimpleKGPipeline:
         llm (LLMInterface): An instance of an LLM to use for entity and relation extraction.
         driver (neo4j.Driver): A Neo4j driver instance for database connection.
         embedder (Embedder): An instance of an embedder used to generate chunk embeddings from text chunks.
-        entities (Optional[List[Union[str, dict[str, str], SchemaEntity]]]): A list of either:
+        schema (Optional[Union[SchemaConfig, dict[str, list]]]): A schema configuration defining entities,
+                                                   relations, and potential schema relationships.
+                                                   This is the recommended way to provide schema information.
+        entities (Optional[List[Union[str, dict[str, str], SchemaEntity]]]): DEPRECATED. A list of either:
 
             - str: entity labels
             - dict: following the SchemaEntity schema, ie with label, description and properties keys
 
-        relations (Optional[List[Union[str, dict[str, str], SchemaRelation]]]): A list of either:
+        relations (Optional[List[Union[str, dict[str, str], SchemaRelation]]]): DEPRECATED. A list of either:
 
             - str: relation label
             - dict: following the SchemaRelation schema, ie with label, description and properties keys
 
-        potential_schema (Optional[List[tuple]]): A list of potential schema relationships.
+        potential_schema (Optional[List[tuple]]): DEPRECATED. A list of potential schema relationships.
         enforce_schema (str): Validation of the extracted entities/rels against the provided schema. Defaults to "NONE", where schema enforcement will be ignored even if the schema is provided. Possible values "None" or "STRICT".
         from_pdf (bool): Determines whether to include the PdfLoader in the pipeline.
                          If True, expects `file_path` input in `run` methods.
@@ -85,6 +92,7 @@ def __init__(
         entities: Optional[Sequence[EntityInputType]] = None,
         relations: Optional[Sequence[RelationInputType]] = None,
         potential_schema: Optional[List[tuple[str, str, str]]] = None,
+        schema: Optional[Union[SchemaConfig, dict[str, list]]] = None,
         enforce_schema: str = "NONE",
         from_pdf: bool = True,
         text_splitter: Optional[TextSplitter] = None,
@@ -96,15 +104,65 @@ def __init__(
         lexical_graph_config: Optional[LexicalGraphConfig] = None,
         neo4j_database: Optional[str] = None,
     ):
+        # deprecation warnings for old parameters
+        if any([entities, relations, potential_schema]) and schema is not None:
+            logger.warning(
+                "Both 'schema' and individual schema components (entities, relations, potential_schema) "
+                "were provided. The 'schema' parameter takes precedence. In the future, individual "
+                "components will be removed. Please use only the 'schema' parameter."
+            )
+            # emit a DeprecationWarning for tools that might be monitoring for it
+            warnings.warn(
+                "Both 'schema' and individual schema components are provided. Use only 'schema'.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        elif any([entities, relations, potential_schema]):
+            logger.warning(
+                "The 'entities', 'relations', and 'potential_schema' parameters are deprecated "
+                "and will be removed in a future version. "
+                "Please use the 'schema' parameter instead."
+            )
+            warnings.warn(
+                "The 'entities', 'relations', and 'potential_schema' parameters are deprecated.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            
+        # handle schema precedence over individual schema components
+        schema_entities = []
+        schema_relations = []
+        schema_potential = None
+        
+        if schema is not None:
+            # schema takes precedence over individual components
+            if isinstance(schema, SchemaConfig):
+                # use the SchemaConfig directly
+                pass
+            else:
+                # convert dictionary to entity/relation lists
+                schema_entities = schema.get("entities", [])
+                schema_relations = schema.get("relations", [])
+                schema_potential = schema.get("potential_schema")
+        else:
+            # Use the individual components if provided
+            schema_entities = entities or []
+            schema_relations = relations or []
+            schema_potential = potential_schema
+        
+        # determine if automatic schema extraction should be performed
+        has_schema = bool(schema_entities or schema_relations or schema_potential or isinstance(schema, SchemaConfig))
+        auto_schema_extraction = not has_schema
+            
         try:
             config = SimpleKGPipelineConfig(
                 # argument type are fixed in the Config object
                 llm_config=llm,  # type: ignore[arg-type]
                 neo4j_config=driver,  # type: ignore[arg-type]
                 embedder_config=embedder,  # type: ignore[arg-type]
-                entities=entities or [],
-                relations=relations or [],
-                potential_schema=potential_schema,
+                entities=schema_entities,
+                relations=schema_relations,
+                potential_schema=schema_potential,
                 enforce_schema=SchemaEnforcementMode(enforce_schema),
                 from_pdf=from_pdf,
                 pdf_loader=ComponentType(pdf_loader) if pdf_loader else None,
@@ -115,6 +173,7 @@ def __init__(
                 perform_entity_resolution=perform_entity_resolution,
                 lexical_graph_config=lexical_graph_config,
                 neo4j_database=neo4j_database,
+                auto_schema_extraction=auto_schema_extraction,
             )
         except (ValidationError, ValueError) as e:
             raise PipelineDefinitionError() from e