From d5d1070e67886cb993330637194e5d892bc3e235 Mon Sep 17 00:00:00 2001 From: Estelle Scifo Date: Tue, 8 Oct 2024 11:32:13 +0200 Subject: [PATCH] Change field value from "CONTINUE" to "IGNORE" to match enum field name (#171) * Change field value from "IGNORE" to "CONTINUE" to match enum field name * Update CHANGELOG.md --- CHANGELOG.md | 3 +++ .../experimental/components/entity_relation_extractor.py | 6 +++++- src/neo4j_graphrag/experimental/pipeline/kg_builder.py | 6 +++--- tests/unit/experimental/pipeline/test_kg_builder.py | 6 +++--- 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 59ad9fccb..c56a091ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,9 @@ ### Fixed - Fix a bug where `openai` Python client and `numpy` were required to import any embedder or LLM. +### Changed +- The value associated to the enum field `OnError.IGNORE` has been changed from "CONTINUE" to "IGNORE" to stick to the convention and match the field name. + ## 1.0.0a1 ## 1.0.0a0 diff --git a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py index 7d712e0a2..b84f60275 100644 --- a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py +++ b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py @@ -45,7 +45,11 @@ class OnError(enum.Enum): RAISE = "RAISE" - IGNORE = "CONTINUE" + IGNORE = "IGNORE" + + @classmethod + def possible_values(cls) -> List[str]: + return [e.value for e in cls] CHUNK_NODE_LABEL = "Chunk" diff --git a/src/neo4j_graphrag/experimental/pipeline/kg_builder.py b/src/neo4j_graphrag/experimental/pipeline/kg_builder.py index 96d6476fe..9f7ff488e 100644 --- a/src/neo4j_graphrag/experimental/pipeline/kg_builder.py +++ b/src/neo4j_graphrag/experimental/pipeline/kg_builder.py @@ -81,7 +81,7 @@ class SimpleKGPipeline: text_splitter (Optional[Any]): A text splitter component. Defaults to FixedSizeSplitter(). pdf_loader (Optional[Any]): A PDF loader component. Defaults to PdfLoader(). kg_writer (Optional[Any]): A knowledge graph writer component. Defaults to Neo4jWriter(). - on_error (str): Error handling strategy. Defaults to "CONTINUE". Possible values: "RAISE" or "CONTINUE". + on_error (str): Error handling strategy for the Entity and relation extractor. Defaults to "IGNORE", where chunk will be ignored if extraction fails. Possible values: "RAISE" or "IGNORE". perform_entity_resolution (bool): Merge entities with same label and name. Default: True prompt_template (str): A custom prompt template to use for extraction. """ @@ -98,7 +98,7 @@ def __init__( text_splitter: Optional[Any] = None, pdf_loader: Optional[Any] = None, kg_writer: Optional[Any] = None, - on_error: str = "CONTINUE", + on_error: str = "IGNORE", prompt_template: Union[ERExtractionTemplate, str] = ERExtractionTemplate(), perform_entity_resolution: bool = True, ): @@ -110,7 +110,7 @@ def __init__( on_error_enum = OnError(on_error) except ValueError: raise PipelineDefinitionError( - f"Invalid value for on_error: {on_error}. Expected 'RAISE' or 'CONTINUE'." + f"Invalid value for on_error: {on_error}. Expected one of {OnError.possible_values()}." ) config = SimpleKGPipelineConfig( diff --git a/tests/unit/experimental/pipeline/test_kg_builder.py b/tests/unit/experimental/pipeline/test_kg_builder.py index 8624c3937..188180f29 100644 --- a/tests/unit/experimental/pipeline/test_kg_builder.py +++ b/tests/unit/experimental/pipeline/test_kg_builder.py @@ -259,10 +259,10 @@ def test_simple_kg_pipeline_on_error_invalid_value() -> None: llm=llm, driver=driver, embedder=embedder, - on_error="IGNORE", + on_error="INVALID_VALUE", ) - assert "Expected 'RAISE' or 'CONTINUE'" in str(exc_info.value) + assert "Expected one of ['RAISE', 'IGNORE']" in str(exc_info.value) def test_simple_kg_pipeline_no_entity_resolution() -> None: @@ -274,7 +274,7 @@ def test_simple_kg_pipeline_no_entity_resolution() -> None: llm=llm, driver=driver, embedder=embedder, - on_error="CONTINUE", + on_error="IGNORE", perform_entity_resolution=False, )