refactor: Refactor SplitText component to SplitContentComponent

This commit is contained in:
Rodrigo 2024-06-15 22:47:35 -03:00
commit 4f0f8f3e0e
2 changed files with 58 additions and 39 deletions

View file

@ -1,49 +1,67 @@
from typing import Optional from typing import List
from langflow.custom import CustomComponent from langflow.custom import Component
from langflow.field_typing import Text from langflow.inputs import IntInput, StrInput, HandleInput
from langflow.schema import Data from langflow.schema import Data
from langflow.template import Output
from langflow.utils.util import unescape_string from langflow.utils.util import unescape_string
class SplitTextComponent(CustomComponent): class SplitContentComponent(Component):
display_name: str = "Split Text" display_name: str = "Split Content"
description: str = "Split text into chunks of a specified length." description: str = "Split textual content into chunks of a specified length."
icon = "split"
def build_config(self): inputs = [
return { HandleInput(
"inputs": { name="data",
"display_name": "Inputs", display_name="Data",
"info": "Texts to split.", info="Data with text to split.",
"input_types": ["Data", "Text"], input_types=["Data"]
}, ),
"separator": { StrInput(
"display_name": "Separator", name="content_key",
"info": 'The character to split on. Defaults to " ".', display_name="Content Key",
}, info="The key to access the text content in the Data object.",
"truncate_size": { value="content",
"display_name": "Truncate Size", ),
"info": "The maximum length (in number of characters) of each chunk to keep. Defaults to 0 (no truncation).", StrInput(
}, name="separator",
} display_name="Separator",
info='The character to split on. Defaults to "\n".',
value="\n",
advanced=True
),
IntInput(
name="chunk_size",
display_name="Chunk Size",
info="The maximum length (in number of characters) of each chunk. Defaults to 0 (no chunking).",
value=0,
advanced=True
),
]
def build( outputs = [
self, Output(display_name="Chunks", name="chunks", method="split_text"),
inputs: list[Text], ]
separator: str = " ",
truncate_size: Optional[int] = 0,
) -> list[Data]:
separator = unescape_string(separator)
outputs = [] def split_text(self) -> List[Data]:
for text in inputs: data = self.data if isinstance(self.data, list) else [self.data]
chunks = text.split(separator) content_key = self.content_key
separator = unescape_string(self.separator)
chunk_size = self.chunk_size
results = []
if truncate_size: for row in data:
chunks = [chunk[:truncate_size] for chunk in chunks] content = row.data.get(content_key, '')
if chunk_size > 0:
chunks = [content[i:i + chunk_size] for i in range(0, len(content), chunk_size)]
else:
chunks = content.split(separator)
for chunk in chunks: for chunk in chunks:
outputs.append(Data(data={"parent": text, "text": chunk})) if chunk.strip():
results.append(Data(data={"parent": content, "text": chunk}))
self.status = outputs self.status = results
return outputs return results

View file

@ -16,14 +16,15 @@ __all__ = [
"ClearMessageHistoryComponent", "ClearMessageHistoryComponent",
"ExtractKeyFromDataComponent", "ExtractKeyFromDataComponent",
"FlowToolComponent", "FlowToolComponent",
"ListFlowsComponent",
"ListenComponent", "ListenComponent",
"ListFlowsComponent",
"MergeDataComponent", "MergeDataComponent",
"NotifyComponent", "NotifyComponent",
"PythonFunctionComponent", "PythonFunctionComponent",
"PythonFunctionComponent",
"RunFlowComponent", "RunFlowComponent",
"RunnableExecComponent", "RunnableExecComponent",
"SQLExecutorComponent", "SQLExecutorComponent",
"SubFlowComponent", "SubFlowComponent",
"PythonFunctionComponent", "SplitContentComponent",
] ]