Coverage for src / lexigram / contracts / ai / document.py: 93%

1"""Document and Text Splitter contracts for Lexigram. 

2 

3Defines Document and text splitting classes analogous to LangChain's. 

4""" 

5 

6from __future__ import annotations 

7 

8from dataclasses import dataclass, field 

9from typing import Any 

10 

11 

12@dataclass(frozen=True) 

13class Document: 

14 """A document (like LangChain's Document). 

15 

16 Attributes: 

17 page_content: The text content. 

18 metadata: Additional metadata. 

19 """ 

20 

21 page_content: str 

22 metadata: dict[str, Any] = field(default_factory=dict) 

23 

24 

25class TextSplitter: 

26 """Base text splitter (like LangChain's TextSplitter). 

27 

28 Splits text into chunks. 

29 """ 

30 

31 def __init__( 

32 self, 

33 chunk_size: int = 1000, 

34 chunk_overlap: int = 200, 

35 ) -> None: 

36 self.chunk_size = chunk_size 

37 self.chunk_overlap = chunk_overlap 

38 

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

40 """Split text into chunks. 

41 

42 Args: 

43 text: Text to split. 

44 

45 Returns: 

46 List of text chunks. 

47 """ 

48 if not text: 

49 return [] 

50 

51 chunks = [] 

52 start = 0 

53 while start < len(text): 

54 end = start + self.chunk_size 

55 chunk = text[start:end] 

56 chunks.append(chunk) 

57 start += self.chunk_size - self.chunk_overlap 

58 

59 return chunks 

60 

61 

62class RecursiveCharacterTextSplitter(TextSplitter): 

63 """Split by characters recursively (like LangChain's RecursiveCharacterTextSplitter).""" 

64 

65 def __init__( 

66 self, 

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

68 chunk_size: int = 1000, 

69 chunk_overlap: int = 200, 

70 ) -> None: 

71 super().__init__(chunk_size=chunk_size, chunk_overlap=chunk_overlap) 

72 self.separators = separators or ["\n\n", "\n", " ", ""] 

73 

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

75 """Split text recursively by separators.""" 

76 if not text: 

77 return [] 

78 

79 final_chunks: list[str] = [] 

80 for separator in self.separators: 

81 chunks = [] 

82 if separator == "": 

83 chunks = list(text) 

84 else: 

85 parts = text.split(separator) 

86 for part in parts: 

87 if part: 

88 chunks.append(part) 

89 

90 if len(chunks) > 1: 

91 break 

92 

93 if not chunks: 

94 return [text] 

95 

96 return self._merge_chunks(chunks) 

97 

98 def _merge_chunks(self, chunks: list[str]) -> list[str]: 

99 """Merge chunks respecting size limits.""" 

100 merged = [] 

101 current = "" 

102 

103 for chunk in chunks: 

104 if len(current) + len(chunk) <= self.chunk_size: 

105 current += chunk 

106 else: 

107 if current: 

108 merged.append(current) 

109 current = chunk 

110 

111 if current: 

112 merged.append(current) 

113 

114 return merged 

115 

116 

117__all__ = [ 

118 "Document", 

119 "RecursiveCharacterTextSplitter", 

120 "TextSplitter", 

121]