1"""Semantic chunking strategy."""
2
3from __future__ import annotations
4
5import re
6from typing import Any
7
8from lexigram.ai.rag.chunking.base import AbstractChunker
9from lexigram.ai.rag.chunking.types import Chunk
10
11
12class SemanticChunker(AbstractChunker):
13 """Semantic chunking based on sentence/paragraph boundaries.
14
15 Splits text at natural boundaries (sentences, paragraphs) while respecting
16 size constraints.
17
18 Example:
19 >>> chunker = SemanticChunker(max_chunk_size=1000)
20 >>> chunks = chunker.chunk("Document with sentences...")
21 """
22
23 def __init__(
24 self,
25 max_chunk_size: int = 1000,
26 min_chunk_size: int = 100,
27 prefer_paragraphs: bool = True,
28 ):
29 """Initialize semantic chunker.
30
31 Args:
32 max_chunk_size: Maximum chunk size in characters
33 min_chunk_size: Minimum chunk size (to avoid too-small chunks)
34 prefer_paragraphs: Prefer paragraph boundaries over sentences
35 """
36 self.max_chunk_size = max_chunk_size
37 self.min_chunk_size = min_chunk_size
38 self.prefer_paragraphs = prefer_paragraphs
39
40 # Sentence boundary pattern
41 self.sentence_pattern = re.compile(r"(?<=[.!?])\s+")
42 # Paragraph boundary pattern
43 self.paragraph_pattern = re.compile(r"\n\n+")
44
45 def chunk(self, text: str, metadata: dict[str, Any] | None = None) -> list[Chunk]:
46 """Split text semantically.
47
48 Args:
49 text: Text to chunk
50 metadata: Optional metadata
51
52 Returns:
53 List of chunks
54 """
55 if not text:
56 return []
57
58 chunks: list[Chunk] = []
59 chunk_index = 0
60 current_chunk: list[str] = []
61 current_size: int = 0
62 chunk_start: int = 0
63
64 if self.prefer_paragraphs:
65 # Split by paragraphs
66 paragraphs = self.paragraph_pattern.split(text)
67 current_chunk = []
68 current_size = 0
69 chunk_start = 0
70
71 for raw_para in paragraphs:
72 para = raw_para.strip()
73 if not para:
74 continue
75
76 # If this paragraph alone exceeds max size, split it by sentences
77 if len(para) > self.max_chunk_size:
78 # Save current chunk if any
79 if current_chunk and current_size >= self.min_chunk_size:
80 chunk_text = "\n\n".join(current_chunk)
81 chunks.append(
82 Chunk(
83 text=chunk_text,
84 start_index=chunk_start,
85 end_index=chunk_start + len(chunk_text),
86 chunk_index=chunk_index,
87 metadata=metadata,
88 ),
89 )
90 chunk_index += 1
91
92 # Split large paragraph by sentences
93 sentences = self.sentence_pattern.split(para)
94 for raw_sent in sentences:
95 sent = raw_sent.strip()
96 if sent and len(sent) >= self.min_chunk_size:
97 chunks.append(
98 Chunk(
99 text=sent,
100 start_index=text.find(sent),
101 end_index=text.find(sent) + len(sent),
102 chunk_index=chunk_index,
103 metadata=metadata,
104 ),
105 )
106 chunk_index += 1
107
108 current_chunk = []
109 current_size = 0
110 chunk_start = text.find(para) + len(para)
111 continue
112
113 # Check if adding this paragraph would exceed size
114 if current_size + len(para) > self.max_chunk_size and current_chunk:
115 # Save current chunk
116 if current_size >= self.min_chunk_size:
117 chunk_text = "\n\n".join(current_chunk)
118 chunks.append(
119 Chunk(
120 text=chunk_text,
121 start_index=chunk_start,
122 end_index=chunk_start + len(chunk_text),
123 chunk_index=chunk_index,
124 metadata=metadata,
125 ),
126 )
127 chunk_index += 1
128
129 current_chunk = [para]
130 current_size = len(para)
131 chunk_start = text.find(para)
132 else:
133 current_chunk.append(para)
134 current_size += len(para) + 2 # +2 for \n\n
135
136 # Add remaining chunk
137 if current_chunk and current_size >= self.min_chunk_size:
138 chunk_text = "\n\n".join(current_chunk)
139 chunks.append(
140 Chunk(
141 text=chunk_text,
142 start_index=chunk_start,
143 end_index=chunk_start + len(chunk_text),
144 chunk_index=chunk_index,
145 metadata=metadata,
146 ),
147 )
148
149 else:
150 # Split by sentences
151 sentences = self.sentence_pattern.split(text)
152 current_chunk = []
153 current_size = 0
154 chunk_start = 0
155
156 for raw_sent in sentences:
157 sent = raw_sent.strip()
158 if not sent:
159 continue
160
161 if current_size + len(sent) > self.max_chunk_size and current_chunk:
162 if current_size >= self.min_chunk_size:
163 chunk_text = " ".join(current_chunk)
164 chunks.append(
165 Chunk(
166 text=chunk_text,
167 start_index=chunk_start,
168 end_index=chunk_start + len(chunk_text),
169 chunk_index=chunk_index,
170 metadata=metadata,
171 ),
172 )
173 chunk_index += 1
174
175 current_chunk = [sent]
176 current_size = len(sent)
177 chunk_start = text.find(sent)
178 else:
179 current_chunk.append(sent)
180 current_size += len(sent) + 1 # +1 for space
181
182 # Add remaining
183 if current_chunk and current_size >= self.min_chunk_size:
184 chunk_text = " ".join(current_chunk)
185 chunks.append(
186 Chunk(
187 text=chunk_text,
188 start_index=chunk_start,
189 end_index=chunk_start + len(chunk_text),
190 chunk_index=chunk_index,
191 metadata=metadata,
192 ),
193 )
194
195 return chunks