klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[OpenSearch] Add chunks specific to an app_id if present (#765)

Browse Source
This commit is contained in:
Deshraj Yadav
2023-10-04 15:46:22 -07:00
committed by GitHub
parent 352e71461d
commit 64a34cac32
6 changed files with 81 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/chunkers/test_image_chunker.py
View File

@@ -17,14 +17,15 @@ class TestImageChunker(unittest.TestCase):
chunker.set_data_type(DataType.IMAGES)
image_path = "./tmp/image.jpeg"
result = chunker.create_chunks(MockLoader(), image_path)
app_id = "app1"
result = chunker.create_chunks(MockLoader(), image_path, app_id=app_id)
expected_chunks = {
"doc_id": "123",
"doc_id": f"{app_id}--123",
"documents": [image_path],
"embeddings": ["embedding"],
"ids": ["140bedbf9c3f6d56a9846d2ba7088798683f4da0c248231336e6a05679e4fdfe"],
"metadatas": [{"data_type": "images", "doc_id": "123", "url": "none"}],
"metadatas": [{"data_type": "images", "doc_id": f"{app_id}--123", "url": "none"}],
}
self.assertEqual(expected_chunks, result)
@@ -37,14 +38,15 @@ class TestImageChunker(unittest.TestCase):
chunker.set_data_type(DataType.IMAGES)
image_path = "./tmp/image.jpeg"
result = chunker.create_chunks(MockLoader(), image_path)
app_id = "app1"
result = chunker.create_chunks(MockLoader(), image_path, app_id=app_id)
expected_chunks = {
"doc_id": "123",
"doc_id": f"{app_id}--123",
"documents": [image_path],
"embeddings": ["embedding"],
"ids": ["140bedbf9c3f6d56a9846d2ba7088798683f4da0c248231336e6a05679e4fdfe"],
"metadatas": [{"data_type": "images", "doc_id": "123", "url": "none"}],
"metadatas": [{"data_type": "images", "doc_id": f"{app_id}--123", "url": "none"}],
}
self.assertEqual(expected_chunks, result)

38
tests/chunkers/test_text.py
View File

@@ -1,31 +1,35 @@
# ruff: noqa: E501
import unittest
from embedchain.chunkers.text import TextChunker
from embedchain.config import ChunkerConfig
from embedchain.models.data_type import DataType
class TestTextChunker(unittest.TestCase):
def test_chunks(self):
class TestTextChunker:
def test_chunks_without_app_id(self):
"""
Test the chunks generated by TextChunker.
# TODO: Not a very precise test.
"""
chunker_config = ChunkerConfig(chunk_size=10, chunk_overlap=0, length_function=len)
chunker = TextChunker(config=chunker_config)
text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
# Data type must be set manually in the test
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
documents = result["documents"]
assert len(documents) > 5
self.assertGreaterEqual(len(documents), 5)
# Additional test cases can be added to cover different scenarios
def test_chunks_with_app_id(self):
"""
Test the chunks generated by TextChunker with app_id
"""
chunker_config = ChunkerConfig(chunk_size=10, chunk_overlap=0, length_function=len)
chunker = TextChunker(config=chunker_config)
text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
documents = result["documents"]
assert len(documents) > 5
def test_big_chunksize(self):
"""
@@ -36,12 +40,9 @@ class TestTextChunker(unittest.TestCase):
text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
# Data type must be set manually in the test
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
documents = result["documents"]
self.assertEqual(len(documents), 1)
assert len(documents) == 1
def test_small_chunksize(self):
"""
@@ -53,14 +54,9 @@ class TestTextChunker(unittest.TestCase):
text = """0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~ \t\n\r\x0b\x0c"""
# Data type must be set manually in the test
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
documents = result["documents"]
print(documents)
self.assertEqual(len(documents), len(text))
assert len(documents) == len(text)
def test_word_count(self):
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len)
@@ -69,7 +65,7 @@ class TestTextChunker(unittest.TestCase):
document = ["ab cd", "ef gh"]
result = chunker.get_word_count(document)
self.assertEqual(result, 4)
assert result == 4
class MockLoader: