llama-index

1
from typing import List, Optional
2

3
from llama_index.legacy.bridge.pydantic import Field
4
from llama_index.legacy.postprocessor.types import BaseNodePostprocessor
5
from llama_index.legacy.schema import MetadataMode, NodeWithScore, QueryBundle
6

7

8
class MetadataReplacementPostProcessor(BaseNodePostprocessor):
9
    target_metadata_key: str = Field(
10
        description="Target metadata key to replace node content with."
11
    )
12

13
    def __init__(self, target_metadata_key: str) -> None:
14
        super().__init__(target_metadata_key=target_metadata_key)
15

16
    @classmethod
17
    def class_name(cls) -> str:
18
        return "MetadataReplacementPostProcessor"
19

20
    def _postprocess_nodes(
21
        self,
22
        nodes: List[NodeWithScore],
23
        query_bundle: Optional[QueryBundle] = None,
24
    ) -> List[NodeWithScore]:
25
        for n in nodes:
26
            n.node.set_content(
27
                n.node.metadata.get(
28
                    self.target_metadata_key,
29
                    n.node.get_content(metadata_mode=MetadataMode.NONE),
30
                )
31
            )
32

33
        return nodes
34