22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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 | class NodeParser(TransformComponent, ABC):
"""Base interface for node parser."""
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=CallbackManager, exclude=True
)
id_func: Callable = Field(
default=None,
description="Function to generate node IDs.",
exclude=True,
)
class Config:
arbitrary_types_allowed = True
@validator("id_func", pre=True)
def _validate_id_func(cls, v: Any) -> Any:
if v is None:
return default_id_func
return v
@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, None)
if parent_doc is not None:
start_char_idx = parent_doc.text.find(
node.get_content(metadata_mode=MetadataMode.NONE)
)
# update start/end char idx
if start_char_idx >= 0:
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:
node.metadata.update(parent_doc.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
):
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
):
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: List[BaseNode], **kwargs: Any) -> List[BaseNode]:
return self.get_nodes_from_documents(nodes, **kwargs)
async def acall(self, nodes: List[BaseNode], **kwargs: Any) -> List[BaseNode]:
return await self.aget_nodes_from_documents(nodes, **kwargs)
|