50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196 | class NodeParser(TransformComponent, ABC):
"""Base interface for node parser."""
model_config = ConfigDict(arbitrary_types_allowed=True)
include_metadata: bool = Field(
default=True, description="Whether or not to consider metadata when splitting."
)
include_prev_next_rel: bool = Field(
default=True, description="Include prev/next node relationships."
)
callback_manager: CallbackManager = Field(
default_factory=lambda: CallbackManager([]), exclude=True
)
id_func: Optional[IdFuncCallable] = Field(
default=None,
description="Function to generate node IDs.",
)
@abstractmethod
def _parse_nodes(
self,
nodes: Sequence[BaseNode],
show_progress: bool = False,
**kwargs: Any,
) -> List[BaseNode]:
...
async def _aparse_nodes(
self,
nodes: Sequence[BaseNode],
show_progress: bool = False,
**kwargs: Any,
) -> List[BaseNode]:
return self._parse_nodes(nodes, show_progress=show_progress, **kwargs)
def _postprocess_parsed_nodes(
self, nodes: List[BaseNode], parent_doc_map: Dict[str, Document]
) -> List[BaseNode]:
for i, node in enumerate(nodes):
parent_doc = parent_doc_map.get(node.ref_doc_id or "", None)
parent_node = node.source_node
if parent_doc is not None:
if parent_doc.source_node is not None:
node.relationships.update(
{
NodeRelationship.SOURCE: parent_doc.source_node,
}
)
start_char_idx = parent_doc.text.find(
node.get_content(metadata_mode=MetadataMode.NONE)
)
# update start/end char idx
if start_char_idx >= 0 and isinstance(node, TextNode):
node.start_char_idx = start_char_idx
node.end_char_idx = start_char_idx + len(
node.get_content(metadata_mode=MetadataMode.NONE)
)
# update metadata
if self.include_metadata:
# Merge parent_doc.metadata into nodes.metadata, giving preference to node's values
node.metadata = {**parent_doc.metadata, **node.metadata}
if parent_node is not None:
if self.include_metadata:
parent_metadata = parent_node.metadata
combined_metadata = {**parent_metadata, **node.metadata}
# Merge parent_node.metadata into nodes.metadata, giving preference to node's values
node.metadata.update(combined_metadata)
if self.include_prev_next_rel:
# establish prev/next relationships if nodes share the same source_node
if (
i > 0
and node.source_node
and nodes[i - 1].source_node
and nodes[i - 1].source_node.node_id == node.source_node.node_id # type: ignore
):
node.relationships[NodeRelationship.PREVIOUS] = nodes[
i - 1
].as_related_node_info()
if (
i < len(nodes) - 1
and node.source_node
and nodes[i + 1].source_node
and nodes[i + 1].source_node.node_id == node.source_node.node_id # type: ignore
):
node.relationships[NodeRelationship.NEXT] = nodes[
i + 1
].as_related_node_info()
return nodes
def get_nodes_from_documents(
self,
documents: Sequence[Document],
show_progress: bool = False,
**kwargs: Any,
) -> List[BaseNode]:
"""Parse documents into nodes.
Args:
documents (Sequence[Document]): documents to parse
show_progress (bool): whether to show progress bar
"""
doc_id_to_document = {doc.id_: doc for doc in documents}
with self.callback_manager.event(
CBEventType.NODE_PARSING, payload={EventPayload.DOCUMENTS: documents}
) as event:
nodes = self._parse_nodes(documents, show_progress=show_progress, **kwargs)
nodes = self._postprocess_parsed_nodes(nodes, doc_id_to_document)
event.on_end({EventPayload.NODES: nodes})
return nodes
async def aget_nodes_from_documents(
self,
documents: Sequence[Document],
show_progress: bool = False,
**kwargs: Any,
) -> List[BaseNode]:
doc_id_to_document = {doc.id_: doc for doc in documents}
with self.callback_manager.event(
CBEventType.NODE_PARSING, payload={EventPayload.DOCUMENTS: documents}
) as event:
nodes = await self._aparse_nodes(
documents, show_progress=show_progress, **kwargs
)
nodes = self._postprocess_parsed_nodes(nodes, doc_id_to_document)
event.on_end({EventPayload.NODES: nodes})
return nodes
def __call__(self, nodes: Sequence[BaseNode], **kwargs: Any) -> List[BaseNode]:
return self.get_nodes_from_documents(nodes, **kwargs) # type: ignore
async def acall(self, nodes: Sequence[BaseNode], **kwargs: Any) -> List[BaseNode]:
return await self.aget_nodes_from_documents(nodes, **kwargs) # type: ignore
|