100 lines
2.7 KiB
Python
100 lines
2.7 KiB
Python
|
# from unittest.mock import Mock, patch
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from guardrails.utils.docs_utils import ( # sentence_split,
|
||
|
|
TextSplitter,
|
||
|
|
get_chunks_from_text,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class MockTokenizer:
|
||
|
|
def encode(self, text):
|
||
|
|
return [token_id for token_id in range(1, len(text) + 1)]
|
||
|
|
|
||
|
|
def decode(self, tokens):
|
||
|
|
return " ".join([str(token_id) for token_id in tokens])
|
||
|
|
|
||
|
|
|
||
|
|
class MockPromptTemplate:
|
||
|
|
def get_prompt_variables(self):
|
||
|
|
return ["var1", "var2"]
|
||
|
|
|
||
|
|
def format(self, **kwargs):
|
||
|
|
return " ".join([f"{key}:{value}" for key, value in kwargs.items()])
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_tokenizer(monkeypatch):
|
||
|
|
mock = MockTokenizer()
|
||
|
|
monkeypatch.setattr("tiktoken.get_encoding", lambda _: mock)
|
||
|
|
return mock
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_prompt_template():
|
||
|
|
return MockPromptTemplate()
|
||
|
|
|
||
|
|
|
||
|
|
def test_text_splitter_split(mock_tokenizer):
|
||
|
|
text_splitter = TextSplitter()
|
||
|
|
text = "This is a test text."
|
||
|
|
chunks = text_splitter.split(text, tokens_per_chunk=10, token_overlap=5, buffer=2)
|
||
|
|
|
||
|
|
assert len(chunks) == 7
|
||
|
|
assert chunks[0] == "1 2 3 4 5 6 7 8"
|
||
|
|
assert chunks[1] == "4 5 6 7 8 9 10 11"
|
||
|
|
assert chunks[2] == "7 8 9 10 11 12 13 14"
|
||
|
|
assert chunks[3] == "10 11 12 13 14 15 16 17"
|
||
|
|
|
||
|
|
|
||
|
|
def test_prompt_template_token_length(mock_tokenizer, mock_prompt_template):
|
||
|
|
text_splitter = TextSplitter()
|
||
|
|
length = text_splitter.prompt_template_token_length(mock_prompt_template)
|
||
|
|
assert length == 11 # Assuming the encoded tokens count is 11
|
||
|
|
|
||
|
|
|
||
|
|
def test_text_splitter_callable(mock_tokenizer):
|
||
|
|
text_splitter = TextSplitter()
|
||
|
|
text = "This is a test text."
|
||
|
|
chunks = text_splitter(text, tokens_per_chunk=10, token_overlap=5, buffer=2)
|
||
|
|
|
||
|
|
assert len(chunks) == 7
|
||
|
|
assert chunks[0] == "1 2 3 4 5 6 7 8"
|
||
|
|
assert chunks[1] == "4 5 6 7 8 9 10 11"
|
||
|
|
assert chunks[2] == "7 8 9 10 11 12 13 14"
|
||
|
|
assert chunks[3] == "10 11 12 13 14 15 16 17"
|
||
|
|
|
||
|
|
|
||
|
|
class MockNLTK:
|
||
|
|
@staticmethod
|
||
|
|
def sent_tokenize(text):
|
||
|
|
return ["sentence1", "sentence2", "sentence3"]
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def word_tokenize(text):
|
||
|
|
return ["word1", "word2", "word3"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_get_chunks_from_text_char():
|
||
|
|
text = "This is a test."
|
||
|
|
chunks = get_chunks_from_text(
|
||
|
|
text, chunk_strategy="char", chunk_size=4, chunk_overlap=1
|
||
|
|
)
|
||
|
|
assert len(chunks) == 5
|
||
|
|
assert chunks[0] == "T h i s"
|
||
|
|
assert chunks[1] == "s i s"
|
||
|
|
assert chunks[2] == "s a "
|
||
|
|
assert chunks[3] == " t e s"
|
||
|
|
assert chunks[4] == "s t ."
|
||
|
|
|
||
|
|
|
||
|
|
def test_get_chunks_from_text_invalid_strategy():
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
get_chunks_from_text(
|
||
|
|
"Invalid strategy.",
|
||
|
|
chunk_strategy="invalid_strategy",
|
||
|
|
chunk_size=1,
|
||
|
|
chunk_overlap=0,
|
||
|
|
)
|