refactor: Update RecursiveCharacterTextSplitterComponent to use type hinting for separators and chunk_size

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-20 18:46:58 -03:00
commit 1b89ac2d73

View file

@ -56,17 +56,17 @@ class RecursiveCharacterTextSplitterComponent(Component):
""" """
if self.separators == "": if self.separators == "":
self.separators = None self.separators: list[str] | None = None
elif self.separators: elif self.separators:
# check if the separators list has escaped characters # check if the separators list has escaped characters
# if there are escaped characters, unescape them # if there are escaped characters, unescape them
self.separators = [unescape_string(x) for x in self.separators] self.separators: list[str] | None = [unescape_string(x) for x in self.separators]
# Make sure chunk_size and chunk_overlap are ints # Make sure chunk_size and chunk_overlap are ints
if self.chunk_size: if self.chunk_size:
self.chunk_size = int(self.chunk_size) self.chunk_size: int = int(self.chunk_size)
if self.chunk_overlap: if self.chunk_overlap:
self.chunk_overlap = int(self.chunk_overlap) self.chunk_overlap: int = int(self.chunk_overlap)
splitter = RecursiveCharacterTextSplitter( splitter = RecursiveCharacterTextSplitter(
separators=self.separators, separators=self.separators,
chunk_size=self.chunk_size, chunk_size=self.chunk_size,
@ -74,7 +74,7 @@ class RecursiveCharacterTextSplitterComponent(Component):
) )
documents = [] documents = []
if not isinstance(self.data_input, list): if not isinstance(self.data_input, list):
self.data_input = [self.data_input] self.data_input: list[Data] = [self.data_input]
for _input in self.data_input: for _input in self.data_input:
if isinstance(_input, Data): if isinstance(_input, Data):
documents.append(_input.to_lc_document()) documents.append(_input.to_lc_document())