Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/chunking/strategies/token.py: 24%

1"""Token-based chunking strategy.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any 

6 

7from lexigram.ai.rag.chunking.base import AbstractChunker 

8from lexigram.ai.rag.chunking.types import Chunk 

9 

10 

11class TokenChunker(AbstractChunker): 

12 """Chunking based on token count using tiktoken. 

13 

14 Example: 

15 >>> chunker = TokenChunker(chunk_size=512, overlap=50) 

16 >>> chunks = chunker.chunk("Long text to tokenize...") 

17 """ 

18 

19 def __init__( 

20 self, 

21 chunk_size: int = 512, 

22 overlap: int = 50, 

23 encoding_name: str = "cl100k_base", 

24 ): 

25 """Initialize token chunker. 

26 

27 Args: 

28 chunk_size: Target tokens per chunk 

29 overlap: Token overlap between chunks 

30 encoding_name: tiktoken encoding name 

31 """ 

32 self.chunk_size = chunk_size 

33 self.overlap = overlap 

34 self.encoding_name = encoding_name 

35 

36 try: 

37 import tiktoken # type: ignore[import-not-found] 

38 

39 self.encoding = tiktoken.get_encoding(encoding_name) 

40 except ImportError as e: 

41 raise ImportError( 

42 "Token chunking requires 'tiktoken' package. " 

43 "Install with: pip install tiktoken", 

44 ) from e 

45 

46 def chunk(self, text: str, metadata: dict[str, Any] | None = None) -> list[Chunk]: 

47 """Split text by tokens. 

48 

49 Args: 

50 text: Text to chunk 

51 metadata: Optional metadata 

52 

53 Returns: 

54 List of chunks 

55 """ 

56 if not text: 

57 return [] 

58 

59 # Encode text to tokens 

60 tokens = self.encoding.encode(text) 

61 

62 # Split into chunks 

63 chunks: list[Chunk] = [] 

64 start = 0 

65 chunk_index = 0 

66 

67 while start < len(tokens): 

68 # Get chunk tokens 

69 end = start + self.chunk_size 

70 chunk_tokens = tokens[start:end] 

71 

72 # Decode tokens back to text for the chunk 

73 chunk_text = self.encoding.decode(chunk_tokens) 

74 

75 # Note: start_index and end_index are character-based in the Chunk model. 

76 # Calculating precise character indices from tokens is complex with tiktoken. 

77 # For now, we'll store basic identifiers. 

78 # If precise mapping is required, more advanced logic would be needed. 

79 

80 chunks.append( 

81 Chunk( 

82 text=chunk_text, 

83 source=metadata.get("source", "unknown") if metadata else "unknown", 

84 chunk_index=chunk_index, 

85 metadata={ 

86 **(metadata or {}), 

87 "token_count": len(chunk_tokens), 

88 "encoding": self.encoding_name, 

89 }, 

90 ), 

91 ) 

92 

93 chunk_index += 1 

94 # Move to next chunk with overlap 

95 # If chunk is shorter than overlap, we still move forward by at least 1 

96 start = max(start + 1, end - self.overlap) 

97 

98 return chunks