jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor SplitTextComponent class to accept inputs of type Record and Text

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-04-04 00:22:37 -03:00
commit 5cade5dca8
1 changed files with 11 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

23
src/backend/base/langflow/components/helpers/SplitText.py
View file

@ -1,14 +1,11 @@
from typing import Optional from typing import Optional
from langchain.text_splitter import ( from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter
RecursiveCharacterTextSplitter,
CharacterTextSplitter,
)
from langchain_core.documents import Document from langchain_core.documents import Document
from langflow.field_typing import Text
from langflow.interface.custom.custom_component import CustomComponent from langflow.interface.custom.custom_component import CustomComponent
from langflow.schema import Record from langflow.schema import Record
from langflow.field_typing import Text
from langflow.utils.util import unescape_string from langflow.utils.util import unescape_string
@ -18,10 +15,10 @@ class SplitTextComponent(CustomComponent):
def build_config(self): def build_config(self):
return { return {
"texts": { "inputs": {
"display_name": "Texts", "display_name": "Inputs",
"info": "Texts to split.", "info": "Texts to split.",
"input_types": ["Text"], "input_types": ["Record", "Text"],
}, },
"separators": { "separators": {
"display_name": "Separators", "display_name": "Separators",
@ -48,7 +45,7 @@ class SplitTextComponent(CustomComponent):
def build( def build(
self, self,
texts: list[Text], inputs: list[Text],
separators: Optional[list[str]] = [" "], separators: Optional[list[str]] = [" "],
chunk_size: Optional[int] = 1000, chunk_size: Optional[int] = 1000,
chunk_overlap: Optional[int] = 200, chunk_overlap: Optional[int] = 200,
@ -77,9 +74,11 @@ class SplitTextComponent(CustomComponent):
) )
documents = [] documents = []
for _text in texts: for _input in inputs:
# documents.append(_input.to_lc_document()) if isinstance(_input, Record):
documents.append(Document(page_content=_text)) documents.append(_input.to_lc_document())
else:
documents.append(Document(page_content=_input))
records = self.to_records(splitter.split_documents(documents)) records = self.to_records(splitter.split_documents(documents))
self.status = records self.status = records