feat: allow caller-supplied metadata on extracted entity nodes and relationships - #590
Open
Aryan-Pardeshi wants to merge 1 commit into
Open
feat: allow caller-supplied metadata on extracted entity nodes and relationships#590Aryan-Pardeshi wants to merge 1 commit into
Aryan-Pardeshi wants to merge 1 commit into
Conversation
…lationships Adds an entity_metadata parameter to SimpleKGPipeline.run_async that is threaded through the pipeline to LLMEntityRelationExtractor. Caller-supplied metadata is merged into every extracted node and relationship property at construction, taking precedence over LLM-produced keys (tenant, owner, or ingestion-time markers). Closes neo4j#588.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #588
The issue lays out the asymmetry precisely, so I will not restate it beyond the summary: document nodes take caller metadata via
document_metadata, chunk nodes take it viaTextChunk.metadata, and extracted entity nodes and relationships have no equivalent — their properties come only from LLM output. The lexical bookkeeping layer is enrichable while the entity layer, the part the library exists to produce, is not.This adds an
entity_metadataparameter threaded fromSimpleKGPipeline.run_async()throughSimpleKGPipelineConfigto the extractor, wherepost_process_chunk()merges it into the properties of every extracted node and relationship.Caller-supplied keys take precedence over LLM-produced ones. That is documented on the parameter: if the model happens to emit a
tenant_id, the caller's value is the one that lands, since the whole point is an authoritative application-level marker.EntityRelationExtractor.run()gains the parameter on the abstract signature, so the example custom extractor inexamples/customize/build_graph/components/extractors/custom_extractor.pyis updated to match. The parameter is optional throughout and nothing is passed when it isNone, so existing pipelines are unaffected.Verified the new tests fail against unmodified
main(2 failed) and pass with the change.tests/unit/experimental/pipeline/test_kg_builder.pyandtests/unit/components/test_entity_relation_extractor.pytogether: 44 passed.mypyclean on both touched source files,ruff checkclean on every file in the diff.One note on the wider suite:
tests/unitshows 27 failures on this branch, butmainshows the same 27 (they come from amistralaiSDK mismatch in my environment), and this branch adds 2 passes on top.If you would rather this lived on a pluggable writer than on the pipeline signature, I am happy to rework it — the reporter sketched a
Neo4jWritersubclass as the current workaround, and I can see the argument either way.