Coverage for src/qdrant_loader/core/chunking/strategy/base_strategy.py: 96%
166 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Base abstract class for chunking strategies."""
3from abc import ABC, abstractmethod
4from typing import TYPE_CHECKING
6import tiktoken
8from qdrant_loader.core.document import Document
9from qdrant_loader.core.text_processing.text_processor import TextProcessor
10from qdrant_loader.utils.logging import LoggingConfig
12if TYPE_CHECKING:
13 from qdrant_loader.config import Settings
15logger = LoggingConfig.get_logger(__name__)
18class BaseChunkingStrategy(ABC):
19 """Base abstract class for all chunking strategies.
21 This class defines the interface that all chunking strategies must implement.
22 Each strategy should provide its own implementation of how to split documents
23 into chunks while preserving their semantic meaning and structure.
24 """
26 # Whether this strategy uses the base TextProcessor (spaCy NER + POS) via
27 # _process_text/_apply_nlp. Strategies that do their own enrichment (docling)
28 # set this False so __init__ does not load a spaCy model they never use.
29 _uses_base_text_processor: bool = True
31 def __init__(
32 self,
33 settings: "Settings",
34 chunk_size: int | None = None,
35 chunk_overlap: int | None = None,
36 ):
37 """Initialize the chunking strategy.
39 Args:
40 settings: Application settings containing configuration for the strategy
41 chunk_size: Maximum number of tokens per chunk (optional, defaults to settings value)
42 chunk_overlap: Number of tokens to overlap between chunks (optional, defaults to settings value)
43 """
44 self.settings = settings
45 self.logger = LoggingConfig.get_logger(self.__class__.__name__)
47 # Initialize token-based chunking parameters
48 self.chunk_size = chunk_size or settings.global_config.chunking.chunk_size
49 self.chunk_overlap = (
50 chunk_overlap or settings.global_config.chunking.chunk_overlap
51 )
52 self.tokenizer = settings.global_config.embedding.tokenizer
54 # Initialize tokenizer based on configuration
55 if self.tokenizer == "none":
56 self.encoding = None
57 else:
58 try:
59 self.encoding = tiktoken.get_encoding(self.tokenizer)
60 except Exception as e:
61 logger.warning(
62 "Failed to initialize tokenizer, falling back to simple character counting",
63 error=str(e),
64 tokenizer=self.tokenizer,
65 )
66 self.encoding = None
68 if self.chunk_overlap >= self.chunk_size:
69 raise ValueError("Chunk overlap must be less than chunk size")
71 # Master switch for NLP metadata extraction across all strategies.
72 self._semantic_analysis_enabled = bool(
73 getattr(settings.global_config.chunking, "enable_semantic_analysis", True)
74 )
76 # Initialize text processor only when NLP metadata extraction is enabled.
77 self.text_processor = (
78 TextProcessor(settings)
79 if self._semantic_analysis_enabled and self._uses_base_text_processor
80 else None
81 )
83 def _count_tokens(self, text: str) -> int:
84 """Count the number of tokens in a text string."""
85 if self.encoding is None:
86 # Fallback to character count if no tokenizer is available
87 return len(text)
88 return len(self.encoding.encode(text))
90 def _process_text(self, text: str) -> dict:
91 """Process text using the text processor.
93 Args:
94 text: Text to process
96 Returns:
97 dict: Processed text features
98 """
99 if self.text_processor is None:
100 return {"tokens": [], "entities": [], "pos_tags": [], "chunks": []}
102 return self.text_processor.process_text(text)
104 def _should_apply_nlp(
105 self, content: str, file_path: str = "", content_type: str = ""
106 ) -> bool:
107 """Determine if NLP processing should be applied to content.
109 Args:
110 content: The content to analyze
111 file_path: File path for extension-based detection
112 content_type: Content type if available
114 Returns:
115 bool: True if NLP processing would be valuable
116 """
117 # Skip NLP for very large content (performance)
118 if len(content) > 20000: # 20KB limit
119 return False
121 # Get file extension
122 ext = ""
123 if file_path and "." in file_path:
124 ext = f".{file_path.lower().split('.')[-1]}"
126 # Skip NLP for code files (except comments/docstrings)
127 code_extensions = {
128 ".py",
129 ".pyx",
130 ".pyi",
131 ".java",
132 ".js",
133 ".jsx",
134 ".mjs",
135 ".ts",
136 ".tsx",
137 ".go",
138 ".rs",
139 ".cpp",
140 ".cc",
141 ".cxx",
142 ".c",
143 ".h",
144 ".cs",
145 ".php",
146 ".rb",
147 ".kt",
148 ".scala",
149 ".swift",
150 ".dart",
151 ".sh",
152 ".bash",
153 ".zsh",
154 ".sql",
155 ".r",
156 ".m",
157 ".pl",
158 ".lua",
159 ".vim",
160 ".asm",
161 }
162 if ext in code_extensions:
163 return False
165 # Skip NLP for structured data files
166 structured_extensions = {
167 ".json",
168 ".xml",
169 ".yaml",
170 ".yml",
171 ".toml",
172 ".ini",
173 ".cfg",
174 ".conf",
175 ".csv",
176 ".tsv",
177 ".log",
178 ".properties",
179 }
180 if ext in structured_extensions:
181 return False
183 # Skip NLP for binary/encoded content
184 binary_extensions = {
185 ".pdf",
186 ".docx",
187 ".xls",
188 ".xlsx",
189 ".pptx",
190 ".zip",
191 ".tar",
192 ".gz",
193 ".bz2",
194 ".7z",
195 ".rar",
196 ".jpg",
197 ".jpeg",
198 ".png",
199 ".gif",
200 ".bmp",
201 ".svg",
202 ".mp3",
203 ".mp4",
204 ".avi",
205 ".mov",
206 ".wav",
207 ".flac",
208 }
209 if ext in binary_extensions:
210 return False
212 # Apply NLP for documentation and text files
213 text_extensions = {".md", ".txt", ".rst", ".adoc", ".tex", ".rtf"}
214 if ext in text_extensions:
215 return True
217 # Apply NLP for HTML content (but be selective)
218 if ext in {".html", ".htm"} or content_type == "html":
219 return True
221 # For unknown extensions, check content characteristics
222 if not ext:
223 # Skip if content looks like code (high ratio of special characters)
224 special_chars = sum(1 for c in content if c in "{}[]();,=<>!&|+-*/%^~`")
225 if len(content) > 0 and special_chars / len(content) > 0.15:
226 return False
228 # Skip if content looks like structured data
229 if content.strip().startswith(("{", "[", "<")) or "=" in content[:100]:
230 return False
232 # Default to applying NLP for text-like content
233 return True
235 def _extract_nlp_worthy_content(self, content: str, element_type: str = "") -> str:
236 """Extract only the parts of content that are worth NLP processing.
238 For code files, this extracts comments and docstrings.
239 For other files, returns the full content.
241 Args:
242 content: The content to process
243 element_type: Type of code element (if applicable)
245 Returns:
246 str: Content suitable for NLP processing
247 """
248 # For code elements, only process comments and docstrings
249 if element_type in ["comment", "docstring"]:
250 return content
252 # For other code elements, extract comments
253 if element_type in ["function", "method", "class", "module"]:
254 return self._extract_comments_and_docstrings(content)
256 # For non-code content, return as-is
257 return content
259 def _extract_comments_and_docstrings(self, code_content: str) -> str:
260 """Extract comments and docstrings from code content.
262 Args:
263 code_content: Code content to extract from
265 Returns:
266 str: Extracted comments and docstrings
267 """
268 extracted_text = []
269 lines = code_content.split("\n")
271 in_multiline_comment = False
272 in_docstring = False
273 docstring_delimiter = None
275 for line in lines:
276 stripped = line.strip()
278 # Python/Shell style comments
279 if stripped.startswith("#"):
280 comment = stripped[1:].strip()
281 if comment: # Skip empty comments
282 extracted_text.append(comment)
284 # C/Java/JS style single line comments
285 elif "//" in stripped:
286 comment_start = stripped.find("//")
287 comment = stripped[comment_start + 2 :].strip()
288 if comment:
289 extracted_text.append(comment)
291 # C/Java/JS style multiline comments
292 elif "/*" in stripped and not in_multiline_comment:
293 in_multiline_comment = True
294 comment_start = stripped.find("/*")
295 comment = stripped[comment_start + 2 :]
296 if "*/" in comment:
297 comment = comment[: comment.find("*/")]
298 in_multiline_comment = False
299 comment = comment.strip()
300 if comment:
301 extracted_text.append(comment)
303 elif in_multiline_comment:
304 if "*/" in stripped:
305 comment = stripped[: stripped.find("*/")]
306 in_multiline_comment = False
307 else:
308 comment = stripped
309 comment = comment.strip("* \t")
310 if comment:
311 extracted_text.append(comment)
313 # Python docstrings
314 elif ('"""' in stripped or "'''" in stripped) and not in_docstring:
315 for delimiter in ['"""', "'''"]:
316 if delimiter in stripped:
317 in_docstring = True
318 docstring_delimiter = delimiter
319 start_idx = stripped.find(delimiter)
320 docstring_content = stripped[start_idx + 3 :]
322 # Check if docstring ends on same line
323 if delimiter in docstring_content:
324 end_idx = docstring_content.find(delimiter)
325 docstring_text = docstring_content[:end_idx].strip()
326 if docstring_text:
327 extracted_text.append(docstring_text)
328 in_docstring = False
329 docstring_delimiter = None
330 else:
331 if docstring_content.strip():
332 extracted_text.append(docstring_content.strip())
333 break
335 elif in_docstring and docstring_delimiter:
336 if docstring_delimiter in stripped:
337 end_idx = stripped.find(docstring_delimiter)
338 docstring_text = stripped[:end_idx].strip()
339 if docstring_text:
340 extracted_text.append(docstring_text)
341 in_docstring = False
342 docstring_delimiter = None
343 else:
344 if stripped:
345 extracted_text.append(stripped)
347 return "\n".join(extracted_text)
349 def _create_chunk_document(
350 self,
351 original_doc: Document,
352 chunk_content: str,
353 chunk_index: int,
354 total_chunks: int,
355 skip_nlp: bool = False,
356 ) -> Document:
357 """Create a new document for a chunk with enhanced metadata.
359 Args:
360 original_doc: Original document
361 chunk_content: Content of the chunk
362 chunk_index: Index of the chunk
363 total_chunks: Total number of chunks
364 skip_nlp: Whether to skip expensive NLP processing
366 Returns:
367 Document: New document instance for the chunk
368 """
369 # Create enhanced metadata
370 metadata = original_doc.metadata.copy()
371 metadata.update(
372 {
373 "chunk_index": chunk_index,
374 "total_chunks": total_chunks,
375 }
376 )
378 # Smart NLP decision based on content type and characteristics
379 file_path = original_doc.metadata.get("file_name", "") or original_doc.source
380 content_type = original_doc.content_type or ""
381 element_type = metadata.get("element_type", "")
383 # For converted files, use the converted content type instead of original file extension
384 conversion_method = metadata.get("conversion_method")
385 if conversion_method == "markitdown":
386 # File was converted to markdown, so treat it as markdown for NLP purposes
387 file_path = "converted.md" # Use .md extension for NLP decision
388 content_type = "md"
390 nlp_applicable = self._should_apply_nlp(chunk_content, file_path, content_type)
392 should_apply_nlp = (
393 self._semantic_analysis_enabled
394 and not skip_nlp
395 and len(chunk_content) <= 10000 # Size limit
396 and total_chunks <= 50 # Chunk count limit
397 and nlp_applicable
398 )
400 if not should_apply_nlp:
401 # Skip NLP processing
402 skip_reason = "performance_optimization"
403 if not self._semantic_analysis_enabled:
404 skip_reason = "semantic_analysis_disabled"
405 elif len(chunk_content) > 10000:
406 skip_reason = "chunk_too_large"
407 elif total_chunks > 50:
408 skip_reason = "too_many_chunks"
409 elif not nlp_applicable:
410 skip_reason = "content_type_inappropriate"
412 metadata.update(
413 {
414 "entities": [],
415 "pos_tags": [],
416 "nlp_skipped": True,
417 "skip_reason": skip_reason,
418 }
419 )
420 else:
421 try:
422 # For code content, only process comments/docstrings
423 nlp_content = self._extract_nlp_worthy_content(
424 chunk_content, element_type
425 )
427 if nlp_content.strip():
428 # Process the NLP-worthy content
429 processed = self._process_text(nlp_content)
430 metadata.update(
431 {
432 "entities": processed["entities"],
433 "pos_tags": processed["pos_tags"],
434 "nlp_skipped": False,
435 "nlp_content_extracted": len(nlp_content)
436 < len(chunk_content),
437 "nlp_content_ratio": (
438 len(nlp_content) / len(chunk_content)
439 if chunk_content
440 else 0
441 ),
442 }
443 )
444 else:
445 # No NLP-worthy content found
446 metadata.update(
447 {
448 "entities": [],
449 "pos_tags": [],
450 "nlp_skipped": True,
451 "skip_reason": "no_nlp_worthy_content",
452 }
453 )
454 except Exception as e:
455 self.logger.warning(
456 f"NLP processing failed for chunk {chunk_index}: {e}"
457 )
458 metadata.update(
459 {
460 "entities": [],
461 "pos_tags": [],
462 "nlp_skipped": True,
463 "skip_reason": "nlp_error",
464 }
465 )
467 return Document(
468 content=chunk_content,
469 metadata=metadata,
470 source=original_doc.source,
471 source_type=original_doc.source_type,
472 url=original_doc.url,
473 title=original_doc.title,
474 content_type=original_doc.content_type,
475 )
477 @abstractmethod
478 def chunk_document(self, document: Document) -> list[Document]:
479 """Split a document into chunks while preserving metadata.
481 This method should:
482 1. Split the document content into appropriate chunks
483 2. Preserve all metadata from the original document
484 3. Add chunk-specific metadata (e.g., chunk index, total chunks)
485 4. Return a list of new Document instances
487 Args:
488 document: The document to chunk
490 Returns:
491 List of chunked documents with preserved metadata
493 Raises:
494 NotImplementedError: If the strategy doesn't implement this method
495 """
496 raise NotImplementedError(
497 "Chunking strategy must implement chunk_document method"
498 )