Siddhant Medar
55643b5a04
Feature/saas hybrid smart chunking * WIP: Initial work on SaaS HSC integration * clean up hsc funcs script * fix merging dict issue * Merge remote-tracking branch 'origin/main' into feature/saas-hybrid-smart-chunking * added changes as per feedback comments * add tests and remove old smart chunking funcs * add more tests Approved-by: Katon Minhas
451 lines
15 KiB
Python
451 lines
15 KiB
Python
"""
|
|
Unit tests for src.investment.hybrid_smart_chunking_funcs module
|
|
|
|
Tests hybrid smart chunking functions including:
|
|
- Context clipping by token count
|
|
- Document preparation for chunking
|
|
- Chunk formatting for LLM
|
|
- Page number extraction
|
|
- Preprocessing pipeline
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
from langchain.schema import Document
|
|
from src.investment.hybrid_smart_chunking_funcs import (
|
|
clip_context_by_tokens,
|
|
prepare_documents_for_chunking,
|
|
format_chunks_for_llm,
|
|
add_page_numbers_to_chunks,
|
|
preprocessing_pipeline,
|
|
)
|
|
|
|
|
|
class TestClipContextByTokens:
|
|
"""Test clip_context_by_tokens function"""
|
|
|
|
def test_context_under_limit(self):
|
|
"""Test context that doesn't need clipping"""
|
|
context = "Short text" * 100 # ~1100 chars = ~314 tokens
|
|
max_tokens = 1000
|
|
|
|
result = clip_context_by_tokens(context, max_tokens)
|
|
|
|
assert result == context
|
|
assert len(result) == len(context)
|
|
|
|
def test_context_over_limit(self):
|
|
"""Test context that needs clipping"""
|
|
context = "A" * 350000 # ~100k tokens
|
|
max_tokens = 10000 # Clip to 10k tokens
|
|
|
|
result = clip_context_by_tokens(context, max_tokens)
|
|
|
|
# Should be clipped to approximately 10k tokens (35k chars)
|
|
assert len(result) < len(context)
|
|
assert len(result) <= int(max_tokens * 3.5)
|
|
|
|
def test_empty_context(self):
|
|
"""Test empty context"""
|
|
result = clip_context_by_tokens("")
|
|
|
|
assert result == ""
|
|
|
|
def test_none_context(self):
|
|
"""Test None context"""
|
|
result = clip_context_by_tokens(None)
|
|
|
|
assert result == ""
|
|
|
|
def test_exact_limit(self):
|
|
"""Test context exactly at the limit"""
|
|
max_tokens = 1000
|
|
context = "A" * int(max_tokens * 3.5) # Exactly at limit
|
|
|
|
result = clip_context_by_tokens(context, max_tokens)
|
|
|
|
assert result == context
|
|
|
|
def test_custom_max_tokens(self):
|
|
"""Test with custom max_tokens parameter"""
|
|
context = "A" * 100000
|
|
max_tokens = 5000
|
|
|
|
result = clip_context_by_tokens(context, max_tokens)
|
|
|
|
estimated_tokens = int(len(result) / 3.5)
|
|
assert estimated_tokens <= max_tokens
|
|
|
|
def test_token_estimation_accuracy(self):
|
|
"""Test that token estimation formula is applied correctly"""
|
|
# 3500 chars should be ~1000 tokens
|
|
context = "A" * 3500
|
|
max_tokens = 1000
|
|
|
|
result = clip_context_by_tokens(context, max_tokens)
|
|
|
|
# Should not clip since it's right at the limit
|
|
assert len(result) == 3500
|
|
|
|
|
|
class TestPrepareDocumentsForChunking:
|
|
"""Test prepare_documents_for_chunking function"""
|
|
|
|
def test_single_section(self):
|
|
"""Test preparing a single section"""
|
|
sections = {
|
|
"Section 1": {
|
|
"content": "Section content here",
|
|
"pages": [1, 2]
|
|
}
|
|
}
|
|
filename = "test_contract.pdf"
|
|
|
|
docs = prepare_documents_for_chunking(sections, filename)
|
|
|
|
assert len(docs) == 1
|
|
assert docs[0].page_content == "Section content here"
|
|
assert docs[0].metadata["section_header"] == "Section 1"
|
|
assert docs[0].metadata["pages"] == [1, 2]
|
|
assert docs[0].metadata["filename"] == filename
|
|
|
|
def test_multiple_sections(self):
|
|
"""Test preparing multiple sections"""
|
|
sections = {
|
|
"Section 1": {"content": "Content 1", "pages": [1]},
|
|
"Section 2": {"content": "Content 2", "pages": [2, 3]},
|
|
"Section 3": {"content": "Content 3", "pages": [4]}
|
|
}
|
|
|
|
docs = prepare_documents_for_chunking(sections)
|
|
|
|
assert len(docs) == 3
|
|
assert all(isinstance(doc, Document) for doc in docs)
|
|
|
|
def test_empty_sections(self):
|
|
"""Test with empty sections dictionary"""
|
|
sections = {}
|
|
|
|
docs = prepare_documents_for_chunking(sections)
|
|
|
|
assert docs == []
|
|
|
|
def test_section_metadata(self):
|
|
"""Test that metadata is correctly set"""
|
|
sections = {
|
|
"Test Section": {
|
|
"content": "This is test content",
|
|
"pages": [5, 6, 7]
|
|
}
|
|
}
|
|
filename = "contract.pdf"
|
|
|
|
docs = prepare_documents_for_chunking(sections, filename)
|
|
|
|
metadata = docs[0].metadata
|
|
assert metadata["section_header"] == "Test Section"
|
|
assert metadata["section_length"] == len("This is test content")
|
|
assert metadata["pages"] == [5, 6, 7]
|
|
assert metadata["filename"] == filename
|
|
|
|
def test_default_filename(self):
|
|
"""Test default filename when not provided"""
|
|
sections = {
|
|
"Section": {"content": "Content", "pages": [1]}
|
|
}
|
|
|
|
docs = prepare_documents_for_chunking(sections)
|
|
|
|
assert docs[0].metadata["filename"] == "contract"
|
|
|
|
|
|
class TestFormatChunksForLLM:
|
|
"""Test format_chunks_for_llm function"""
|
|
|
|
def test_empty_docs(self):
|
|
"""Test with empty document list"""
|
|
text_dict = {"1": "Page 1 content"}
|
|
|
|
result = format_chunks_for_llm(text_dict, [])
|
|
|
|
assert result == ""
|
|
|
|
def test_single_document(self):
|
|
"""Test with single document"""
|
|
text_dict = {
|
|
"1": "Page 1 content",
|
|
"2": "Page 2 content"
|
|
}
|
|
docs = [
|
|
Document(
|
|
page_content="chunk1",
|
|
metadata={"pages": [1]}
|
|
)
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
assert "Page 1 content" in result
|
|
assert "Page 2 content" not in result
|
|
|
|
def test_multiple_documents_different_pages(self):
|
|
"""Test with documents from different pages"""
|
|
text_dict = {
|
|
"1": "Page 1 content",
|
|
"2": "Page 2 content",
|
|
"3": "Page 3 content"
|
|
}
|
|
docs = [
|
|
Document(page_content="chunk1", metadata={"pages": [1]}),
|
|
Document(page_content="chunk2", metadata={"pages": [3]})
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
assert "Page 1 content" in result
|
|
assert "Page 3 content" in result
|
|
assert "Page 2 content" not in result
|
|
|
|
def test_overlapping_pages(self):
|
|
"""Test with documents that reference overlapping pages"""
|
|
text_dict = {
|
|
"1": "Page 1 content",
|
|
"2": "Page 2 content"
|
|
}
|
|
docs = [
|
|
Document(page_content="chunk1", metadata={"pages": [1, 2]}),
|
|
Document(page_content="chunk2", metadata={"pages": [2]})
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
# Both pages should appear only once (set deduplication)
|
|
assert "Page 1 content" in result
|
|
assert "Page 2 content" in result
|
|
|
|
def test_missing_page_in_text_dict(self):
|
|
"""Test when document references page not in text_dict"""
|
|
text_dict = {
|
|
"1": "Page 1 content"
|
|
}
|
|
docs = [
|
|
Document(page_content="chunk", metadata={"pages": [1, 99]})
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
# Should only include page 1, skip page 99
|
|
assert "Page 1 content" in result
|
|
|
|
def test_none_values_in_text_dict(self):
|
|
"""Test handling of None values in text_dict"""
|
|
text_dict = {
|
|
"1": "Page 1 content",
|
|
"2": None,
|
|
"3": "Page 3 content"
|
|
}
|
|
docs = [
|
|
Document(page_content="chunk", metadata={"pages": [1, 2, 3]})
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
# Should skip None value for page 2
|
|
assert "Page 1 content" in result
|
|
assert "Page 3 content" in result
|
|
|
|
def test_document_without_pages_metadata(self):
|
|
"""Test document without pages in metadata"""
|
|
text_dict = {"1": "Page 1 content"}
|
|
docs = [
|
|
Document(page_content="chunk", metadata={})
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
# Should return empty or handle gracefully
|
|
assert isinstance(result, str)
|
|
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.clip_context_by_tokens')
|
|
def test_calls_clip_context(self, mock_clip):
|
|
"""Test that format_chunks_for_llm calls clip_context_by_tokens"""
|
|
mock_clip.return_value = "clipped context"
|
|
text_dict = {"1": "Page 1 content"}
|
|
docs = [Document(page_content="chunk", metadata={"pages": [1]})]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
mock_clip.assert_called_once()
|
|
assert result == "clipped context"
|
|
|
|
def test_integer_page_numbers(self):
|
|
"""Test handling of integer page numbers (converted to strings)"""
|
|
text_dict = {
|
|
"1": "Page 1 content",
|
|
"2": "Page 2 content"
|
|
}
|
|
docs = [
|
|
Document(page_content="chunk", metadata={"pages": [1, 2]}) # Integers
|
|
]
|
|
|
|
result = format_chunks_for_llm(text_dict, docs)
|
|
|
|
assert "Page 1 content" in result
|
|
assert "Page 2 content" in result
|
|
|
|
|
|
class TestAddPageNumbersToChunks:
|
|
"""Test add_page_numbers_to_chunks function"""
|
|
|
|
def test_single_section_with_page_markers(self):
|
|
"""Test extracting page numbers from section with markers"""
|
|
sections = {
|
|
"Section 1": "Start of Page No. = 1\nContent here\nStart of Page No. = 2\nMore content"
|
|
}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
assert "Section 1" in result
|
|
assert result["Section 1"]["content"] == sections["Section 1"]
|
|
assert 1 in result["Section 1"]["pages"]
|
|
assert 2 in result["Section 1"]["pages"]
|
|
|
|
def test_multiple_sections(self):
|
|
"""Test multiple sections with page markers"""
|
|
sections = {
|
|
"Section 1": "Start of Page No. = 5\nContent",
|
|
"Section 2": "Start of Page No. = 6\nMore content\nStart of Page No. = 7\nEven more"
|
|
}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
assert 5 in result["Section 1"]["pages"]
|
|
assert 6 in result["Section 2"]["pages"]
|
|
assert 7 in result["Section 2"]["pages"]
|
|
|
|
def test_section_without_markers(self):
|
|
"""Test section without page markers"""
|
|
sections = {
|
|
"Section 1": "Start of Page No. = 1\nContent",
|
|
"Section 2": "Content without markers"
|
|
}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
# Section 2 should inherit last page from Section 1
|
|
assert result["Section 2"]["pages"] == [1]
|
|
|
|
def test_empty_sections(self):
|
|
"""Test with empty sections dictionary"""
|
|
sections = {}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
assert result == {}
|
|
|
|
def test_page_number_deduplication(self):
|
|
"""Test that duplicate page numbers are removed"""
|
|
sections = {
|
|
"Section 1": "Start of Page No. = 3\nContent\nStart of Page No. = 3\nMore"
|
|
}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
# Should have unique page numbers
|
|
assert result["Section 1"]["pages"].count(3) == 1
|
|
|
|
def test_page_numbers_are_sorted(self):
|
|
"""Test that page numbers are sorted"""
|
|
sections = {
|
|
"Section 1": "Start of Page No. = 5\nStart of Page No. = 2\nStart of Page No. = 8"
|
|
}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
pages = result["Section 1"]["pages"]
|
|
assert pages == sorted(pages)
|
|
|
|
def test_alternative_marker_format(self):
|
|
"""Test alternative page marker format without period"""
|
|
sections = {
|
|
"Section 1": "Start of Page No = 10\nContent"
|
|
}
|
|
|
|
result = add_page_numbers_to_chunks(sections)
|
|
|
|
assert 10 in result["Section 1"]["pages"]
|
|
|
|
|
|
class TestPreprocessingPipeline:
|
|
"""Test preprocessing_pipeline function"""
|
|
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.hybrid_smart_chunking_preprocessing.split_contract_text_into_sections')
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.add_page_numbers_to_chunks')
|
|
def test_pipeline_calls_functions(self, mock_add_pages, mock_split):
|
|
"""Test that pipeline calls the necessary functions"""
|
|
contract_text = "Test contract text"
|
|
|
|
# Mock split_contract_text_into_sections return
|
|
mock_sections = {"Section 1": "content"}
|
|
mock_verification = (True, "Perfect preservation")
|
|
mock_split.return_value = (mock_sections, contract_text, "\n", mock_verification)
|
|
|
|
# Mock add_page_numbers_to_chunks return
|
|
mock_sections_with_pages = {
|
|
"Section 1": {"content": "content", "pages": [1]}
|
|
}
|
|
mock_add_pages.return_value = mock_sections_with_pages
|
|
|
|
sections, is_preserved = preprocessing_pipeline(contract_text)
|
|
|
|
mock_split.assert_called_once_with(contract_text)
|
|
mock_add_pages.assert_called_once()
|
|
assert is_preserved is True
|
|
assert sections == mock_sections_with_pages
|
|
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.hybrid_smart_chunking_preprocessing.split_contract_text_into_sections')
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.add_page_numbers_to_chunks')
|
|
def test_pipeline_returns_preservation_status(self, mock_add_pages, mock_split):
|
|
"""Test that pipeline returns preservation status"""
|
|
# Test successful preservation
|
|
mock_split.return_value = ({}, "text", "\n", (True, "Success"))
|
|
mock_add_pages.return_value = {}
|
|
|
|
sections, is_preserved = preprocessing_pipeline("text")
|
|
assert is_preserved is True
|
|
|
|
# Test failed preservation
|
|
mock_split.return_value = ({}, "text", "\n", (False, "Failed"))
|
|
mock_add_pages.return_value = {}
|
|
|
|
sections, is_preserved = preprocessing_pipeline("text")
|
|
assert is_preserved is False
|
|
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.hybrid_smart_chunking_preprocessing.split_contract_text_into_sections')
|
|
@patch('src.investment.hybrid_smart_chunking_funcs.add_page_numbers_to_chunks')
|
|
def test_pipeline_with_real_text(self, mock_add_pages, mock_split):
|
|
"""Test pipeline with realistic contract text"""
|
|
contract_text = """1 Definitions
|
|
Definition content
|
|
2 Terms
|
|
Terms content"""
|
|
|
|
mock_sections = {
|
|
"1 Definitions": "Definition content",
|
|
"2 Terms": "Terms content"
|
|
}
|
|
mock_split.return_value = (mock_sections, contract_text, "\n", (True, "OK"))
|
|
|
|
mock_sections_with_pages = {
|
|
"1 Definitions": {"content": "Definition content", "pages": [1]},
|
|
"2 Terms": {"content": "Terms content", "pages": [2]}
|
|
}
|
|
mock_add_pages.return_value = mock_sections_with_pages
|
|
|
|
sections, is_preserved = preprocessing_pipeline(contract_text)
|
|
|
|
assert len(sections) == 2
|
|
assert "1 Definitions" in sections
|
|
assert "2 Terms" in sections
|
|
assert is_preserved is True
|