Skip to content

Commit 2bc1018

Browse files
committed
Use BaseSchemaBuilder
1 parent 9379eca commit 2bc1018

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
@@ -44,6 +44,7 @@
4444
SchemaBuilder,
4545
GraphSchema,
4646
SchemaFromTextExtractor,
47+
BaseSchemaBuilder,
4748
)
4849
from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter
4950
from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
@@ -178,7 +179,7 @@ def _get_run_params_for_splitter(self) -> dict[str, Any]:
178179
def _get_chunk_embedder(self) -> TextChunkEmbedder:
179180
return TextChunkEmbedder(embedder=self.get_default_embedder())
180181

181-
def _get_schema(self) -> Union[SchemaBuilder, SchemaFromTextExtractor]:
182+
def _get_schema(self) -> BaseSchemaBuilder:
182183
"""
183184
Get the appropriate schema component based on configuration.
184185
Return SchemaFromTextExtractor for automatic extraction or SchemaBuilder for manual schema.

0 commit comments

Comments
 (0)