참고소스 수정본

참고/guardrails-main/tests/unit_tests/utils/test_docs_utils.py

@@ -0,0 +1,99 @@
+# 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,
+        )