neo4j / neo4j/neo4j-graphrag-python

No way to attach caller-supplied metadata to extracted entity nodes

Open
#588 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.3k
Forks
246
Avg merge
1d 8h
Merged PRs (30d)
11

Description

Summary

SimpleKGPipeline supports attaching arbitrary caller-supplied metadata to document nodes and to chunk nodes, but there is no equivalent mechanism for the extracted entity nodes and relationships. Their properties come solely from the LLM's output.

This makes it impossible, through the supported API, to stamp application-level metadata — a tenant/owner ID, a source system ID, an ingestion timestamp, an ACL tag — onto the part of the graph the library exists to produce.

The asymmetry

Against main (v1.18.0):

Layer Caller metadata? Where
Document nodes run_async(document_metadata=...)LexicalGraphBuilder.create_document_node(), lexical_graph.py:114 (**document_metadata)
Chunk nodes TextChunk.metadatacreate_chunk_node(), lexical_graph.py:144 (chunk_properties.update(chunk.metadata))
Extracted entities — no mechanism

There is also no way to intercept the graph between extraction and writing: _get_extractor() always returns LLMEntityRelationExtractor and _get_pruner() always returns GraphPruning(). Of the pipeline's stages, only file_loader, pdf_loader, kg_writer and text_splitter accept a caller-supplied component.

The net effect is that the lexical layer — the bookkeeping — is enrichable, while the entity layer is not.

Why it matters

Any deployment where one database serves multiple users or customers needs an owner marker on every node in order to scope or filter retrieval. Today that is only achievable by subclassing Neo4jWriter and stamping properties during the write:

class TenantWriter(Neo4jWriter):
    def __init__(self, *args, tenant_id: str, **kwargs):
        super().__init__(*args, **kwargs)
        self.tenant_id = tenant_id

    async def run(self, graph: Neo4jGraph,
                  lexical_graph_config: LexicalGraphConfig = LexicalGraphConfig()
                  ) -> KGWriterModel:
        for node in graph.nodes:
            node.properties["tenant_id"] = self.tenant_id
        for rel in graph.relationships:
            rel.properties["tenant_id"] = self.tenant_id
        return await super().run(graph, lexical_graph_config)

This works, and it needs no fork — kg_writer is a constructor parameter. But it puts enrichment in the persistence layer, and every multi-tenant user ends up writing the same boilerplate.

Proposal

Add an entity_metadata run parameter, an exact sibling of the existing document_metadata:

await kg.run_async(
    file_path="report.pdf",
    document_metadata={"source": "s3://bucket/report.pdf"},
    entity_metadata={"tenant_id": "acme"},
)

Plumbing would mirror document_metadata: run_async()get_run_params()run_params["extractor"]["entity_metadata"], applied in EntityRelationExtractor.post_process_chunk() (entity_relation_extractor.py:288). That point is on the always-executed path, chunk_graph holds exactly the extracted entities there, and update_ids() immediately prior already normalises node.properties to a dict.

Naming: entity_metadata rather than node_metadata, because the target set is precisely the set the writer already marks __Entity__ (kg_writer.py:229) — documents and chunks are Neo4jNodes too, so node_metadata would imply it applies to them.

Semantics I'd suggest, open to direction:

  • Applies to extracted nodes and relationships — an untagged edge between two tenants' nodes still leaks structure
  • Caller-supplied values win on key collision with an LLM-extracted property, so a hallucinated property can't override an owner ID
  • Reserved keys rejected up front with a clear error (Neo4jNode already rejects id as a property name)
  • Lexical nodes untouched

The change would be purely additive — one run-time parameter alongside an existing one, exposing no new components, with no existing behaviour altered.

Note on entities vs node_types

SimpleKGPipeline.entities is deprecated in favour of schema, and SchemaEntity was renamed NodeType. That rename applied to schema input vocabulary. The __Entity__ label, EntityRelationExtractor and perform_entity_resolution remain current, which is why entity_metadata still seems the right name — but happy to follow whatever naming you prefer.


Happy to open a PR with tests, docs and examples if this direction is welcome. Wanted to check the API shape with you first.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with SimpleKGPipeline.run_async() and get_run_params(), then trace the extractor path into EntityRelationExtractor.post_process_chunk() and the related tests. Review lexical_graph.py and kg_writer.py for existing metadata handling and entity labeling. Done means caller-supplied metadata reaches extracted nodes and relationships without changing lexical nodes, with tests, documentation, and examples covering the proposed API semantics.

Written by the indexing model from the issue text.

Assessment

Tech stack
neo4j, python
Domain
backend-api-design, databases
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.