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

[Vasundhhara]

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.metadata → create_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.

[veenus-rajpoot]

i would like to work on this