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

1"""Custom chunking strategy.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import Any 

7 

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

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

10 

11 

12class CustomChunker(AbstractChunker): 

13 """Custom chunking using a user-defined function. 

14 

15 Example: 

16 >>> def my_splitter(text: str) -> list[str]: 

17 ... return text.split("---") 

18 >>> chunker = CustomChunker(split_fn=my_splitter) 

19 >>> chunks = chunker.chunk("Part 1---Part 2---Part 3") 

20 """ 

21 

22 def __init__(self, split_fn: Callable[[str], list[str]]): 

23 """Initialize custom chunker. 

24 

25 Args: 

26 split_fn: Function that takes text and returns list of chunk strings 

27 """ 

28 self.split_fn = split_fn 

29 

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

31 """Split text using custom function. 

32 

33 Args: 

34 text: Text to chunk 

35 metadata: Optional metadata 

36 

37 Returns: 

38 List of chunks 

39 """ 

40 if not text: 

41 return [] 

42 

43 chunk_strings = self.split_fn(text) 

44 chunks: list[Chunk] = [] 

45 offset = 0 

46 

47 for i, chunk_text in enumerate(chunk_strings): 

48 # Find this chunk's position in original text 

49 start = text.find(chunk_text, offset) 

50 if start == -1: 

51 # Chunk not found, use offset 

52 start = offset 

53 

54 end = start + len(chunk_text) 

55 

56 if chunk_text.strip(): 

57 chunks.append( 

58 Chunk( 

59 text=chunk_text, 

60 start_index=start, 

61 end_index=end, 

62 chunk_index=i, 

63 metadata=metadata, 

64 ), 

65 ) 

66 

67 offset = end 

68 

69 return chunks