Skip to content

Commit d9c7ee7

Browse files
committed
Use BaseSchemaBuilder
1 parent 61cd172 commit d9c7ee7

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
NodeType,
4949
RelationshipType,
5050
SchemaFromTextExtractor,
51+
BaseSchemaBuilder,
5152
)
5253
from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter
5354
from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
@@ -175,7 +176,7 @@ def _get_run_params_for_splitter(self) -> dict[str, Any]:
175176
def _get_chunk_embedder(self) -> TextChunkEmbedder:
176177
return TextChunkEmbedder(embedder=self.get_default_embedder())
177178

178-
def _get_schema(self) -> Union[SchemaBuilder, SchemaFromTextExtractor]:
179+
def _get_schema(self) -> BaseSchemaBuilder:
179180
"""
180181
Get the appropriate schema component based on configuration.
181182
Return SchemaFromTextExtractor for automatic extraction or SchemaBuilder for manual schema.

0 commit comments

Comments
 (0)