Adds a Text Splitter (#82)

* Added text splitter adapter class

* Added copyright header to new files

* Added __future__ import to text_splitters.py for backwards compatibility of type hints

* Moved text splitter file and tests

* Split text splitter adapter into 2 adapters

* Added optional metadata to text chunks

* Fixed typos

* Moved text splitters inside of the components folder

* Fixed Component import

Author: alexthomas93

poetry.lock

@@ -62,6 +62,7 @@ scipy = [
     {version = "^1", python = "<3.12"},
     {version = "^1.7.0", python = ">=3.12"}
 ]
+llama-index = "^0.10.55"
 pytest-asyncio = "^0.23.8"
 
 [tool.poetry.extras]

pyproject.toml

@@ -0,0 +1,50 @@
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, Optional
+
+from neo4j_genai.pipeline.component import Component, DataModel
+from pydantic import BaseModel
+
+
+class TextChunk(BaseModel):
+    """A chunk of text split from a document by a text splitter.
+
+    Attributes:
+        text (str): The raw chunk text.
+        metadata (Optional[dict[str, Any]]): Metadata associated with this chunk such as the id of the next chunk in the original document.
+    """
+
+    text: str
+    metadata: Optional[dict[str, Any]] = None
+
+
+class TextChunks(DataModel):
+    """A collection of text chunks returned from a text splitter.
+
+    Attributes:
+        chunks (list[TextChunk]): A list of text chunks.
+    """
+
+    chunks: list[TextChunk]
+
+
+class TextSplitter(Component):
+    """Interface for a text splitter."""
+
+    @abstractmethod
+    async def run(self, text: str) -> TextChunks:
+        pass

src/neo4j_genai/components/text_splitters/__init__.py

@@ -0,0 +1,49 @@
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from langchain_text_splitters import TextSplitter as LangChainTextSplitter
+from neo4j_genai.components.text_splitters.base import (
+    TextChunk,
+    TextChunks,
+    TextSplitter,
+)
+
+
+class LangChainTextSplitterAdapter(TextSplitter):
+    """Adapter for LangChain TextSplitters.
+    Allows instances of this class to be used in the knowledge graph builder pipeline.
+
+    Args:
+        text_splitter (LangChainTextSplitter): An instance of LangChain's TextSplitter class.
+    """
+
+    def __init__(self, text_splitter: LangChainTextSplitter) -> None:
+        self.text_splitter = text_splitter
+
+    async def run(self, text: str) -> TextChunks:
+        """
+        Splits text into chunks.
+
+        Args:
+            text (str): The text to split.
+
+        Returns:
+            TextChunks: The text split into chunks.
+        """
+        return TextChunks(
+            chunks=[
+                TextChunk(text=chunk) for chunk in self.text_splitter.split_text(text)
+            ]
+        )

src/neo4j_genai/components/text_splitters/base.py

@@ -0,0 +1,49 @@
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from llama_index.core.node_parser import TextSplitter as LlamaIndexTextSplitter
+from neo4j_genai.components.text_splitters.base import (
+    TextChunk,
+    TextChunks,
+    TextSplitter,
+)
+
+
+class LlamaIndexTextSplitterAdapter(TextSplitter):
+    """Adapter for LlamaIndex TextSplitters.
+    Allows instances of this class to be used in the knowledge graph builder pipeline.
+
+    Args:
+        text_splitter (LlamaIndexTextSplitter): An instance of LlamaIndex's TextSplitter class.
+    """
+
+    def __init__(self, text_splitter: LlamaIndexTextSplitter) -> None:
+        self.text_splitter = text_splitter
+
+    async def run(self, text: str) -> TextChunks:
+        """
+        Splits text into chunks.
+
+        Args:
+            text (str): The text to split.
+
+        Returns:
+            TextChunks: The text split into chunks.
+        """
+        return TextChunks(
+            chunks=[
+                TextChunk(text=chunk) for chunk in self.text_splitter.split_text(text)
+            ]
+        )

src/neo4j_genai/components/text_splitters/langchain.py

@@ -0,0 +1,37 @@
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import pytest
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from neo4j_genai.components.text_splitters.base import TextChunk, TextChunks
+from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter
+
+text = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit.
+In cursus erat quis ornare condimentum. Ut sollicitudin libero nec quam vestibulum, non tristique augue tempor.
+Nulla fringilla, augue ac fermentum ultricies, mauris tellus tempor orci, at tincidunt purus arcu vitae nisl.
+Nunc suscipit neque vitae ipsum viverra, eu interdum tortor iaculis.
+Suspendisse sit amet quam non ipsum molestie euismod finibus eu nisi. Quisque sit amet aliquet leo, vel auctor dolor.
+Sed auctor enim at tempus eleifend. Suspendisse potenti. Suspendisse congue tellus id justo bibendum, at commodo sapien porta.
+Nam sagittis nisl vitae nibh pellentesque, et convallis turpis ultrices.
+"""
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter() -> None:
+    text_splitter = LangChainTextSplitterAdapter(RecursiveCharacterTextSplitter())
+    text_chunks = await text_splitter.run(text)
+    assert isinstance(text_chunks, TextChunks)
+    for text_chunk in text_chunks.chunks:
+        assert isinstance(text_chunk, TextChunk)
+        assert text_chunk.text in text

src/neo4j_genai/components/text_splitters/llamaindex.py

@@ -0,0 +1,39 @@
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import pytest
+from llama_index.core.node_parser.text.sentence import SentenceSplitter
+from neo4j_genai.components.text_splitters.base import TextChunk, TextChunks
+from neo4j_genai.components.text_splitters.llamaindex import (
+    LlamaIndexTextSplitterAdapter,
+)
+
+text = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit.
+In cursus erat quis ornare condimentum. Ut sollicitudin libero nec quam vestibulum, non tristique augue tempor.
+Nulla fringilla, augue ac fermentum ultricies, mauris tellus tempor orci, at tincidunt purus arcu vitae nisl.
+Nunc suscipit neque vitae ipsum viverra, eu interdum tortor iaculis.
+Suspendisse sit amet quam non ipsum molestie euismod finibus eu nisi. Quisque sit amet aliquet leo, vel auctor dolor.
+Sed auctor enim at tempus eleifend. Suspendisse potenti. Suspendisse congue tellus id justo bibendum, at commodo sapien porta.
+Nam sagittis nisl vitae nibh pellentesque, et convallis turpis ultrices.
+"""
+
+
+@pytest.mark.asyncio
+async def test_llamaindex_adapter() -> None:
+    text_splitter = LlamaIndexTextSplitterAdapter(SentenceSplitter())
+    text_chunks = await text_splitter.run(text)
+    assert isinstance(text_chunks, TextChunks)
+    for text_chunk in text_chunks.chunks:
+        assert isinstance(text_chunk, TextChunk)
+        assert text_chunk.text in text