Rename schema from text component

Author: NathalieCharbel

docs/source/api.rst

@@ -77,10 +77,10 @@ SchemaBuilder
 .. autoclass:: neo4j_graphrag.experimental.components.schema.SchemaBuilder
     :members: run
 
-SchemaFromText
-=============
+SchemaFromTextExtractor
+-----------------------
 
-.. autoclass:: neo4j_graphrag.experimental.components.schema.SchemaFromText
+.. autoclass:: neo4j_graphrag.experimental.components.schema.SchemaFromTextExtractor
     :members: run
 
 EntityRelationExtractor

docs/source/user_guide_kg_builder.rst

@@ -139,7 +139,7 @@ This schema information can be provided to the `SimpleKGBuilder` as demonstrated
     )
 
 .. note::
-   By default, if no schema is provided to the SimpleKGPipeline, automatic schema extraction will be performed using the LLM (See the :ref:`Automatic Schema Extraction with SchemaFromText` section.
+   By default, if no schema is provided to the SimpleKGPipeline, automatic schema extraction will be performed using the LLM (See the :ref:`Automatic Schema Extraction with SchemaFromTextExtractor`).
 
 Extra configurations
 --------------------
@@ -817,19 +817,18 @@ Here is a code block illustrating these concepts:
 After validation, this schema is saved in a `SchemaConfig` object, whose dict representation is passed
 to the LLM.
 
-Automatic Schema Extraction with SchemaFromText
-----------------------------------------------
-.. _automatic-schema-extraction:
+Automatic Schema Extraction 
+---------------------------
 
-Instead of manually defining the schema, you can use the `SchemaFromText` component to automatically extract a schema from your text using an LLM:
+Instead of manually defining the schema, you can use the `SchemaFromTextExtractor` component to automatically extract a schema from your text using an LLM:
 
 .. code:: python
 
-    from neo4j_graphrag.experimental.components.schema import SchemaFromText
+    from neo4j_graphrag.experimental.components.schema import SchemaFromTextExtractor
     from neo4j_graphrag.llm import OpenAILLM
 
     # Create the automatic schema extractor
-    schema_extractor = SchemaFromText(
+    schema_extractor = SchemaFromTextExtractor(
         llm=OpenAILLM(
             model_name="gpt-4o",
             model_params={
@@ -839,14 +838,7 @@ Instead of manually defining the schema, you can use the `SchemaFromText` compon
         )
     )
 
-    # Extract schema from text
-    schema_config = await schema_extractor.run(text="Your document text here...")
-
-    # Use the extracted schema with other components
-    extractor = LLMEntityRelationExtractor(llm=llm)
-    result = await extractor.run(chunks=chunks, schema=schema_config)
-
-The `SchemaFromText` component analyzes the text and identifies entity types, relationship types, and their property types. It creates a complete `SchemaConfig` object that can be used in the same way as a manually defined schema.
+The `SchemaFromTextExtractor` component analyzes the text and identifies entity types, relationship types, and their property types. It creates a complete `SchemaConfig` object that can be used in the same way as a manually defined schema.
 
 You can also save and reload the extracted schema:
 

examples/automatic_schema_extraction/schema_from_text.py

@@ -1,7 +1,7 @@
-"""This example demonstrates how to use the SchemaFromText component 
+"""This example demonstrates how to use the SchemaFromTextExtractor component
 to automatically extract a schema from text and save it to JSON and YAML files.
 
-The SchemaFromText component uses an LLM to analyze the text and identify entities,
+The SchemaFromTextExtractor component uses an LLM to analyze the text and identify entities,
 relations, and their properties.
 
 Note: This example requires an OpenAI API key to be set in the .env file.
@@ -12,7 +12,10 @@
 import os
 from dotenv import load_dotenv
 
-from neo4j_graphrag.experimental.components.schema import SchemaFromText, SchemaConfig
+from neo4j_graphrag.experimental.components.schema import (
+    SchemaFromTextExtractor,
+    SchemaConfig,
+)
 from neo4j_graphrag.llm import OpenAILLM
 
 # Load environment variables from .env file
@@ -54,77 +57,81 @@
 """
 
 # Define the file paths for saving the schema
-OUTPUT_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
+OUTPUT_DIR = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data"
+)
 JSON_FILE_PATH = os.path.join(OUTPUT_DIR, "extracted_schema.json")
 YAML_FILE_PATH = os.path.join(OUTPUT_DIR, "extracted_schema.yaml")
 
 
 async def extract_and_save_schema() -> SchemaConfig:
     """Extract schema from text and save it to JSON and YAML files."""
-    
+
     # Define LLM parameters
     llm_model_params = {
         "max_tokens": 2000,
         "response_format": {"type": "json_object"},
         "temperature": 0,  # Lower temperature for more consistent output
     }
-    
+
     # Create the LLM instance
     llm = OpenAILLM(
         model_name="gpt-4o",
         model_params=llm_model_params,
     )
-    
+
     try:
-        # Create a SchemaFromText component with the default template
-        schema_extractor = SchemaFromText(llm=llm)
-        
+        # Create a SchemaFromTextExtractor component with the default template
+        schema_extractor = SchemaFromTextExtractor(llm=llm)
+
         print("Extracting schema from text...")
         # Extract schema from text
         inferred_schema = await schema_extractor.run(text=TEXT)
-        
+
         # Ensure the output directory exists
         os.makedirs(OUTPUT_DIR, exist_ok=True)
-        
+
         print(f"Saving schema to JSON file: {JSON_FILE_PATH}")
         # Save the schema to JSON file
         inferred_schema.store_as_json(JSON_FILE_PATH)
-        
+
         print(f"Saving schema to YAML file: {YAML_FILE_PATH}")
         # Save the schema to YAML file
         inferred_schema.store_as_yaml(YAML_FILE_PATH)
-        
+
         print("\nExtracted Schema Summary:")
         print(f"Entities: {list(inferred_schema.entities.keys())}")
-        print(f"Relations: {list(inferred_schema.relations.keys() if inferred_schema.relations else [])}")
-        
+        print(
+            f"Relations: {list(inferred_schema.relations.keys() if inferred_schema.relations else [])}"
+        )
+
         if inferred_schema.potential_schema:
             print("\nPotential Schema:")
             for entity1, relation, entity2 in inferred_schema.potential_schema:
                 print(f"  {entity1} --[{relation}]--> {entity2}")
-        
+
         return inferred_schema
-    
+
     finally:
         # Close the LLM client
         await llm.async_client.close()
 
 
 async def main() -> None:
     """Run the example."""
-    
+
     # Extract schema and save to files
     schema_config = await extract_and_save_schema()
-    
+
     print(f"\nSchema files have been saved to:")
     print(f"  - JSON: {JSON_FILE_PATH}")
     print(f"  - YAML: {YAML_FILE_PATH}")
-    
+
     print("\nExample of how to load the schema from files:")
     print("  from neo4j_graphrag.experimental.components.schema import SchemaConfig")
     print(f"  schema_from_json = SchemaConfig.from_file('{JSON_FILE_PATH}')")
     print(f"  schema_from_yaml = SchemaConfig.from_file('{YAML_FILE_PATH}')")
 
 
 if __name__ == "__main__":
-    asyncio.run(main()) 
+    asyncio.run(main())

src/neo4j_graphrag/experimental/components/schema.py

@@ -335,7 +335,7 @@ async def run(
         return self.create_schema_model(entities, relations, potential_schema)
 
 
-class SchemaFromText(Component):
+class SchemaFromTextExtractor(Component):
     """
     A component for constructing SchemaConfig objects from the output of an LLM after
     automatic schema extraction from text.

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

@@ -31,10 +31,10 @@
 )
 from neo4j_graphrag.experimental.components.schema import (
     SchemaBuilder,
+    SchemaConfig,
     SchemaEntity,
     SchemaRelation,
-    SchemaFromText,
-    SchemaConfig,
+    SchemaFromTextExtractor,
 )
 from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter
 from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
@@ -164,13 +164,13 @@ 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) -> Union[SchemaBuilder, SchemaFromText]:
+    def _get_schema(self) -> Union[SchemaBuilder, SchemaFromTextExtractor]:
         """
         Get the appropriate schema component based on configuration.
-        Return SchemaFromText for automatic extraction or SchemaBuilder for manual schema.
+        Return SchemaFromTextExtractor 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 SchemaFromTextExtractor(llm=self.get_default_llm())
         return SchemaBuilder()
 
     def _process_schema_with_precedence(

tests/unit/experimental/components/test_schema.py

@@ -24,7 +24,7 @@
     SchemaEntity,
     SchemaProperty,
     SchemaRelation,
-    SchemaFromText,
+    SchemaFromTextExtractor,
     SchemaConfig,
 )
 from pydantic import ValidationError
@@ -498,13 +498,15 @@ def invalid_schema_json() -> str:
 
 
 @pytest.fixture
-def schema_from_text(mock_llm: AsyncMock) -> SchemaFromText:
-    return SchemaFromText(llm=mock_llm)
+def schema_from_text(mock_llm: AsyncMock) -> SchemaFromTextExtractor:
+    return SchemaFromTextExtractor(llm=mock_llm)
 
 
 @pytest.mark.asyncio
 async def test_schema_from_text_run_valid_response(
-    schema_from_text: SchemaFromText, mock_llm: AsyncMock, valid_schema_json: str
+    schema_from_text: SchemaFromTextExtractor,
+    mock_llm: AsyncMock,
+    valid_schema_json: str,
 ) -> None:
     # configure the mock LLM to return a valid schema JSON
     mock_llm.ainvoke.return_value = valid_schema_json
@@ -533,7 +535,9 @@ async def test_schema_from_text_run_valid_response(
 
 @pytest.mark.asyncio
 async def test_schema_from_text_run_invalid_json(
-    schema_from_text: SchemaFromText, mock_llm: AsyncMock, invalid_schema_json: str
+    schema_from_text: SchemaFromTextExtractor,
+    mock_llm: AsyncMock,
+    invalid_schema_json: str,
 ) -> None:
     # configure the mock LLM to return invalid JSON
     mock_llm.ainvoke.return_value = invalid_schema_json
@@ -553,8 +557,10 @@ async def test_schema_from_text_custom_template(
     custom_prompt = "This is a custom prompt with text: {text}"
     custom_template = PromptTemplate(template=custom_prompt, expected_inputs=["text"])
 
-    # create SchemaFromText with the custom template
-    schema_from_text = SchemaFromText(llm=mock_llm, prompt_template=custom_template)
+    # create SchemaFromTextExtractor with the custom template
+    schema_from_text = SchemaFromTextExtractor(
+        llm=mock_llm, prompt_template=custom_template
+    )
 
     # configure mock LLM to return valid JSON and capture the prompt that was sent to it
     mock_llm.ainvoke.return_value = valid_schema_json
@@ -574,8 +580,8 @@ async def test_schema_from_text_llm_params(
     # configure custom LLM parameters
     llm_params = {"temperature": 0.1, "max_tokens": 500}
 
-    # create SchemaFromText with custom LLM parameters
-    schema_from_text = SchemaFromText(llm=mock_llm, llm_params=llm_params)
+    # create SchemaFromTextExtractor with custom LLM parameters
+    schema_from_text = SchemaFromTextExtractor(llm=mock_llm, llm_params=llm_params)
 
     # configure the mock LLM to return a valid schema JSON
     mock_llm.ainvoke.return_value = valid_schema_json