Coverage for src/qdrant_loader/core/text_processing/text_processor.py: 97%

1"""Text processing module integrating LangChain, spaCy, and NLTK.""" 

2 

3import nltk 

4from qdrant_loader.config import Settings 

5from qdrant_loader.core.text_processing import spacy_model_cache 

6from qdrant_loader.utils.logging import LoggingConfig 

7 

8logger = LoggingConfig.get_logger(__name__) 

9 

10# Performance constants to prevent timeouts 

11MAX_TEXT_LENGTH_FOR_SPACY = 100_000 # 100KB limit for spaCy processing 

12MAX_ENTITIES_TO_EXTRACT = 50 # Limit number of entities 

13MAX_POS_TAGS_TO_EXTRACT = 200 # Limit number of POS tags 

14 

15 

16class TextProcessor: 

17 """Text processing service integrating multiple NLP libraries.""" 

18 

19 def __init__(self, settings: Settings): 

20 """Initialize the text processor with required models and configurations. 

21 

22 Args: 

23 settings: Application settings containing configuration for text processing 

24 """ 

25 self.settings = settings 

26 

27 # Download required NLTK data 

28 try: 

29 nltk.data.find("tokenizers/punkt") 

30 except LookupError: 

31 nltk.download("punkt") 

32 try: 

33 nltk.data.find("corpora/stopwords") 

34 except LookupError: 

35 nltk.download("stopwords") 

36 

37 # Load spaCy model with optimized settings. Cached and shared across 

38 # instances -- TextProcessor is constructed fresh per document, and 

39 # spacy.load() is too expensive to repeat for every one of them. 

40 spacy_model = settings.global_config.semantic_analysis.spacy_model 

41 

42 def _load_nlp(): 

43 import spacy 

44 from spacy.cli.download import download 

45 

46 try: 

47 nlp = spacy.load(spacy_model) 

48 except OSError: 

49 logger.info(f"Downloading spaCy model {spacy_model}...") 

50 download(spacy_model) 

51 nlp = spacy.load(spacy_model) 

52 

53 # Optimize spaCy pipeline for speed: keep only essential 

54 # components (tokenizer, tagger, ner), excluding the parser. 

55 if "parser" in nlp.pipe_names: 

56 essential_pipes = [pipe for pipe in nlp.pipe_names if pipe != "parser"] 

57 nlp.select_pipes(enable=essential_pipes) 

58 return nlp 

59 

60 self.nlp = spacy_model_cache.get_or_load( 

61 ("text_processor", spacy_model), _load_nlp 

62 ) 

63 

64 # Initialize LangChain text splitter with configuration from settings 

65 from langchain_text_splitters import RecursiveCharacterTextSplitter 

66 

67 self.text_splitter = RecursiveCharacterTextSplitter( 

68 chunk_size=settings.global_config.chunking.chunk_size, 

69 chunk_overlap=settings.global_config.chunking.chunk_overlap, 

70 length_function=len, 

71 separators=[ 

72 "\n\n", 

73 "\n", 

74 ".", 

75 "!", 

76 "?", 

77 " ", 

78 "", 

79 ], # Added sentence-ending punctuation 

80 ) 

81 

82 def process_text(self, text: str) -> dict: 

83 """Process text using multiple NLP libraries with performance optimizations. 

84 

85 Args: 

86 text: Input text to process 

87 

88 Returns: 

89 dict: Processed text features including: 

90 - tokens: List of tokens (limited) 

91 - entities: List of named entities (limited) 

92 - pos_tags: List of part-of-speech tags (limited) 

93 - chunks: List of text chunks 

94 """ 

95 # Performance check: truncate very long text 

96 if len(text) > MAX_TEXT_LENGTH_FOR_SPACY: 

97 logger.debug( 

98 f"Text too long for spaCy processing ({len(text)} chars), truncating to {MAX_TEXT_LENGTH_FOR_SPACY}" 

99 ) 

100 text = text[:MAX_TEXT_LENGTH_FOR_SPACY] 

101 

102 try: 

103 # Process with spaCy (optimized) 

104 doc = self.nlp(text) 

105 

106 # Extract features with limits to prevent timeouts 

107 tokens = [token.text for token in doc][ 

108 :MAX_POS_TAGS_TO_EXTRACT 

109 ] # Limit tokens 

110 entities = [(ent.text, ent.label_) for ent in doc.ents][ 

111 :MAX_ENTITIES_TO_EXTRACT 

112 ] # Limit entities 

113 pos_tags = [(token.text, token.pos_) for token in doc][ 

114 :MAX_POS_TAGS_TO_EXTRACT 

115 ] # Limit POS tags 

116 

117 # Process with LangChain (fast) 

118 chunks = self.text_splitter.split_text(text) 

119 

120 return { 

121 "tokens": tokens, 

122 "entities": entities, 

123 "pos_tags": pos_tags, 

124 "chunks": chunks, 

125 } 

126 except Exception as e: 

127 logger.warning(f"Text processing failed: {e}") 

128 # Return minimal results on error 

129 return { 

130 "tokens": [], 

131 "entities": [], 

132 "pos_tags": [], 

133 "chunks": [text] if text else [], 

134 } 

135 

136 def get_entities(self, text: str) -> list[tuple]: 

137 """Extract named entities from text using spaCy with performance limits. 

138 

139 Args: 

140 text: Input text 

141 

142 Returns: 

143 List of (entity_text, entity_type) tuples 

144 """ 

145 # Performance check: truncate very long text 

146 if len(text) > MAX_TEXT_LENGTH_FOR_SPACY: 

147 text = text[:MAX_TEXT_LENGTH_FOR_SPACY] 

148 

149 try: 

150 doc = self.nlp(text) 

151 return [(ent.text, ent.label_) for ent in doc.ents][ 

152 :MAX_ENTITIES_TO_EXTRACT 

153 ] 

154 except Exception as e: 

155 logger.warning(f"Entity extraction failed: {e}") 

156 return [] 

157 

158 def get_pos_tags(self, text: str) -> list[tuple]: 

159 """Get part-of-speech tags using spaCy with performance limits. 

160 

161 Args: 

162 text: Input text 

163 

164 Returns: 

165 List of (word, pos_tag) tuples 

166 """ 

167 # Performance check: truncate very long text 

168 if len(text) > MAX_TEXT_LENGTH_FOR_SPACY: 

169 text = text[:MAX_TEXT_LENGTH_FOR_SPACY] 

170 

171 try: 

172 doc = self.nlp(text) 

173 return [(token.text, token.pos_) for token in doc][:MAX_POS_TAGS_TO_EXTRACT] 

174 except Exception as e: 

175 logger.warning(f"POS tagging failed: {e}") 

176 return [] 

177 

178 def split_into_chunks(self, text: str, chunk_size: int | None = None) -> list[str]: 

179 """Split text into chunks using LangChain's text splitter. 

180 

181 Args: 

182 text: Input text 

183 chunk_size: Optional custom chunk size 

184 

185 Returns: 

186 List of text chunks 

187 """ 

188 try: 

189 if chunk_size: 

190 # Create a new text splitter with the custom chunk size 

191 # Ensure chunk_overlap is smaller than chunk_size 

192 from langchain_text_splitters import RecursiveCharacterTextSplitter 

193 

194 chunk_overlap = min(chunk_size // 4, 50) # 25% of chunk size, max 50 

195 text_splitter = RecursiveCharacterTextSplitter( 

196 chunk_size=chunk_size, 

197 chunk_overlap=chunk_overlap, 

198 length_function=len, 

199 separators=[ 

200 "\n\n", 

201 "\n", 

202 ".", 

203 "!", 

204 "?", 

205 " ", 

206 "", 

207 ], # Added sentence-ending punctuation 

208 ) 

209 return text_splitter.split_text(text) 

210 return self.text_splitter.split_text(text) 

211 except Exception as e: 

212 logger.warning(f"Text splitting failed: {e}") 

213 # Return the original text as a single chunk on error 

214 return [text] if text else []