Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/preprocessing/tables.py: 97%

1""" 

2Table extractor for extracting structured tables from documents. 

3 

4Note: This is a simplified implementation. In production, 

5integrate with libraries like camelot, tabula, or deep learning models. 

6""" 

7 

8from __future__ import annotations 

9 

10import re 

11from typing import TYPE_CHECKING 

12 

13if TYPE_CHECKING: 

14 from lexigram.ai.rag.preprocessing.types import ( 

15 DocumentType, 

16 ExtractedTable, 

17 TableFormat, 

18 ) 

19 

20from lexigram.ai.rag.preprocessing.base import AbstractPreprocessor 

21from lexigram.ai.rag.preprocessing.document import PreprocessedDocument 

22from lexigram.ai.rag.preprocessing.types import ( 

23 DocumentMetadata, 

24 DocumentType, 

25 ExtractedTable, 

26 TableFormat, 

27) 

28 

29 

30class TableExtractor(AbstractPreprocessor): 

31 """Table extractor for extracting structured tables from documents. 

32 

33 Note: This is a simplified implementation. In production, 

34 integrate with libraries like camelot, tabula, or deep learning models. 

35 """ 

36 

37 def __init__(self, table_format: TableFormat = TableFormat.MARKDOWN): 

38 """Initialize table extractor. 

39 

40 Args: 

41 table_format: Format for extracted tables. 

42 """ 

43 super().__init__("table_extractor") 

44 self.table_format = table_format 

45 

46 async def preprocess( 

47 self, 

48 content: str, 

49 **kwargs, 

50 ) -> PreprocessedDocument: 

51 """Extract tables from document. 

52 

53 Args: 

54 content: Document content (HTML, Markdown, or text). 

55 **kwargs: Additional parameters. 

56 

57 Returns: 

58 Preprocessed document with extracted tables. 

59 """ 

60 tables = self._extract_tables(content) 

61 

62 # Remove tables from text, replace with references 

63 processed_text = content 

64 for i, table in enumerate(tables): 

65 processed_text = processed_text.replace( 

66 self._get_table_marker(i), 

67 f"[Table {i + 1}: {table.caption or 'Untitled'}]", 

68 ) 

69 

70 metadata = DocumentMetadata( 

71 doc_type=self._detect_document_type(content), 

72 word_count=len(processed_text.split()), 

73 ) 

74 

75 return PreprocessedDocument( 

76 content=processed_text, 

77 metadata=metadata, 

78 tables=tables, 

79 raw_content=content, 

80 ) 

81 

82 def _extract_tables(self, content: str) -> list[ExtractedTable]: 

83 """Extract tables from content. 

84 

85 Args: 

86 content: Document content. 

87 

88 Returns: 

89 List of extracted tables. 

90 """ 

91 tables = [] 

92 

93 # Extract Markdown tables 

94 markdown_tables = self._extract_markdown_tables(content) 

95 tables.extend(markdown_tables) 

96 

97 # Extract HTML tables 

98 html_tables = self._extract_html_tables(content) 

99 tables.extend(html_tables) 

100 

101 return tables 

102 

103 def _extract_markdown_tables(self, content: str) -> list[ExtractedTable]: 

104 """Extract tables from Markdown content. 

105 

106 Args: 

107 content: Markdown content. 

108 

109 Returns: 

110 List of extracted tables. 

111 """ 

112 tables = [] 

113 

114 # Simple Markdown table pattern: | col1 | col2 | 

115 # Followed by separator: |------|------| 

116 # Followed by rows: | val1 | val2 | 

117 

118 # Find table blocks 

119 lines = content.split("\n") 

120 i = 0 

121 table_id = 0 

122 

123 while i < len(lines): 

124 line = lines[i].strip() 

125 

126 # Check if line looks like a table row 

127 if line.startswith("|") and line.endswith("|"): 

128 # Potential table start 

129 table_lines = [line] 

130 i += 1 

131 

132 # Collect all consecutive table lines 

133 while i < len(lines): 

134 next_line = lines[i].strip() 

135 if next_line.startswith("|") and next_line.endswith("|"): 

136 table_lines.append(next_line) 

137 i += 1 

138 else: 

139 break 

140 

141 # Parse table 

142 if len(table_lines) >= 2: # At least header + separator 

143 table = self._parse_markdown_table(table_lines, table_id) 

144 if table: 

145 tables.append(table) 

146 table_id += 1 

147 else: 

148 i += 1 

149 

150 return tables 

151 

152 def _parse_markdown_table( 

153 self, 

154 lines: list[str], 

155 table_id: int, 

156 ) -> ExtractedTable | None: 

157 """Parse Markdown table from lines. 

158 

159 Args: 

160 lines: Table lines. 

161 table_id: Table identifier. 

162 

163 Returns: 

164 Extracted table or None. 

165 """ 

166 if len(lines) < 2: 

167 return None 

168 

169 # Parse headers 

170 header_line = lines[0].strip("|").strip() 

171 headers = list(map(str.strip, header_line.split("|"))) 

172 

173 # Skip separator line (line 1) 

174 # Parse rows 

175 rows = [] 

176 for line in lines[2:]: 

177 row_line = line.strip("|").strip() 

178 row = list(map(str.strip, row_line.split("|"))) 

179 if len(row) == len(headers): # Valid row 

180 rows.append(row) 

181 

182 return ExtractedTable( 

183 table_id=f"table_{table_id}", 

184 page_number=0, # Unknown for text content 

185 headers=headers, 

186 rows=rows, 

187 format=TableFormat.MARKDOWN, 

188 ) 

189 

190 def _extract_html_tables(self, content: str) -> list[ExtractedTable]: 

191 """Extract tables from HTML content. 

192 

193 Args: 

194 content: HTML content. 

195 

196 Returns: 

197 List of extracted tables. 

198 """ 

199 tables = [] 

200 

201 # Simple HTML table extraction using regex 

202 # In production, use proper HTML parser like BeautifulSoup 

203 table_pattern = r"<table[^>]*>(.*?)</table>" 

204 table_matches = re.findall(table_pattern, content, re.DOTALL | re.IGNORECASE) 

205 

206 for i, table_html in enumerate(table_matches): 

207 table = self._parse_html_table(table_html, i) 

208 if table: 

209 tables.append(table) 

210 

211 return tables 

212 

213 def _parse_html_table( 

214 self, 

215 html: str, 

216 table_id: int, 

217 ) -> ExtractedTable | None: 

218 """Parse HTML table. 

219 

220 Args: 

221 html: HTML table content. 

222 table_id: Table identifier. 

223 

224 Returns: 

225 Extracted table or None. 

226 """ 

227 # Extract headers from <th> tags 

228 header_pattern = r"<th[^>]*>(.*?)</th>" 

229 header_matches = re.findall(header_pattern, html, re.DOTALL | re.IGNORECASE) 

230 headers = list(map(self._clean_html, header_matches)) 

231 

232 # Extract rows from <tr> tags 

233 row_pattern = r"<tr[^>]*>(.*?)</tr>" 

234 row_matches = re.findall(row_pattern, html, re.DOTALL | re.IGNORECASE) 

235 

236 rows = [] 

237 for row_html in row_matches: 

238 # Extract cells from <td> tags 

239 cell_pattern = r"<td[^>]*>(.*?)</td>" 

240 cell_matches = re.findall(cell_pattern, row_html, re.DOTALL | re.IGNORECASE) 

241 row = list(map(self._clean_html, cell_matches)) 

242 if row: # Valid row 

243 rows.append(row) 

244 

245 if not headers or not rows: 

246 return None 

247 

248 return ExtractedTable( 

249 table_id=f"table_{table_id}", 

250 page_number=0, # Unknown for HTML content 

251 headers=headers, 

252 rows=rows, 

253 format=self.table_format, 

254 ) 

255 

256 def _clean_html(self, html_text: str) -> str: 

257 """Clean HTML tags from text. 

258 

259 Args: 

260 html_text: Text with HTML tags. 

261 

262 Returns: 

263 Clean text. 

264 """ 

265 # Remove HTML tags 

266 clean = re.sub(r"<[^>]+>", "", html_text) 

267 return clean.strip() 

268 

269 def _detect_document_type(self, content: str) -> DocumentType: 

270 """Detect document type from content. 

271 

272 Args: 

273 content: Document content. 

274 

275 Returns: 

276 Detected document type. 

277 """ 

278 if "<table>" in content.lower(): 

279 return DocumentType.HTML 

280 if "|" in content and "|" in content.split("\n", maxsplit=1)[0]: 

281 return DocumentType.MARKDOWN 

282 return DocumentType.TEXT 

283 

284 def _get_table_marker(self, table_index: int) -> str: 

285 """Get marker for table in content. 

286 

287 Args: 

288 table_index: Index of table. 

289 

290 Returns: 

291 Table marker string. 

292 """ 

293 return f"[TABLE_MARKER_{table_index}]"