Coverage for src/local_deep_research/embeddings/splitters/bounded_semantic_chunker.py: 98%

1"""Bound semantic document chunks with character-based recursive splitters.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Iterable 

6from typing import TYPE_CHECKING, Protocol 

7 

8from langchain_core.documents.transformers import BaseDocumentTransformer 

9from local_deep_research.constants import DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS 

10 

11if TYPE_CHECKING: 

12 from collections.abc import Sequence 

13 

14 from langchain_core.documents import Document 

15 

16 

17class _DocumentSplitter(Protocol): 

18 """Minimal public document-splitting contract.""" 

19 

20 def split_documents( 20 ↛ exitline 20 didn't return from function 'split_documents' because

21 self, documents: Iterable[Document] 

22 ) -> list[Document]: ... 

23 

24 

25class BoundedSemanticChunker(BaseDocumentTransformer): 

26 """Bound semantic-comparison inputs and final chunks by character length.""" 

27 

28 def __init__( 

29 self, 

30 semantic_chunker: _DocumentSplitter, 

31 chunk_size: int, 

32 separators: list[str] | None = None, 

33 ) -> None: 

34 """Create recursive pre- and post-semantic bounding stages. 

35 

36 Note: 

37 The parent ``langchain_text_splitters`` package is imported before 

38 its ``character`` submodule to preserve the parent-first import 

39 invariant required for safe concurrent cold imports. 

40 """ 

41 import langchain_text_splitters # noqa: F401 

42 from langchain_text_splitters.character import ( 

43 RecursiveCharacterTextSplitter, 

44 ) 

45 

46 effective_separators = ( 

47 separators.copy() 

48 if separators is not None 

49 else DEFAULT_LOCAL_SEARCH_TEXT_SEPARATORS.copy() 

50 ) 

51 # RecursiveCharacterTextSplitter only guarantees the hard character 

52 # bound when it can fall back to splitting individual characters. 

53 # User-configured separator lists are allowed to omit that fallback, 

54 # so add it here instead of letting one separator-free span bypass the 

55 # cap this wrapper exists to enforce. 

56 if "" not in effective_separators: 

57 effective_separators.append("") 

58 

59 def create_boundary_splitter() -> RecursiveCharacterTextSplitter: 

60 return RecursiveCharacterTextSplitter( 

61 chunk_size=chunk_size, 

62 # Semantic mode historically did not apply document overlap. 

63 # Zero overlap also prevents a near-size configured overlap 

64 # from amplifying large inputs into excessive requests. 

65 chunk_overlap=0, 

66 length_function=len, 

67 keep_separator=True, 

68 strip_whitespace=False, 

69 separators=effective_separators.copy(), 

70 ) 

71 

72 self._pre_splitter = create_boundary_splitter() 

73 self._semantic_chunker = semantic_chunker 

74 self._post_splitter = create_boundary_splitter() 

75 

76 def split_text(self, text: str) -> list[str]: 

77 return [ 

78 document.page_content for document in self.create_documents([text]) 

79 ] 

80 

81 def create_documents( 

82 self, 

83 texts: list[str], 

84 metadatas: list[dict[str, object]] | None = None, 

85 ) -> list[Document]: 

86 source_documents = self._pre_splitter.create_documents(texts, metadatas) 

87 semantic_documents = self._semantic_chunker.split_documents( 

88 source_documents 

89 ) 

90 return self._post_splitter.split_documents(semantic_documents) 

91 

92 def split_documents(self, documents: Iterable[Document]) -> list[Document]: 

93 """Bound semantic inputs and outputs while preserving their ordering.""" 

94 texts = [] 

95 metadatas = [] 

96 for document in documents: 

97 texts.append(document.page_content) 

98 metadatas.append(document.metadata) 

99 return self.create_documents(texts, metadatas) 

100 

101 def transform_documents( 

102 self, documents: Sequence[Document], **kwargs: object 

103 ) -> Sequence[Document]: 

104 return self.split_documents(documents)