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

1""" 

2Document preprocessing types. 

3 

4Contains dataclasses and enums for document preprocessing. 

5""" 

6 

7from __future__ import annotations 

8 

9from dataclasses import dataclass, field 

10from datetime import datetime 

11from enum import Enum 

12from typing import Any 

13 

14 

15class DocumentType(str, Enum): 

16 """Types of documents that can be preprocessed.""" 

17 

18 PDF = "pdf" 

19 IMAGE = "image" 

20 HTML = "html" 

21 MARKDOWN = "markdown" 

22 TEXT = "text" 

23 DOCX = "docx" 

24 UNKNOWN = "unknown" 

25 

26 

27class TableFormat(str, Enum): 

28 """Formats for extracted tables.""" 

29 

30 MARKDOWN = "markdown" 

31 HTML = "html" 

32 CSV = "csv" 

33 JSON = "json" 

34 

35 

36@dataclass 

37class ExtractedImage: 

38 """Represents an extracted image from a document. 

39 

40 Attributes: 

41 image_id: Unique identifier for the image. 

42 page_number: Page number where image was found. 

43 position: Position on page (x, y, width, height). 

44 alt_text: Alternative text description. 

45 ocr_text: Text extracted from image via OCR. 

46 metadata: Additional metadata. 

47 """ 

48 

49 image_id: str 

50 page_number: int 

51 position: tuple[int, int, int, int] | None = None 

52 alt_text: str | None = None 

53 ocr_text: str | None = None 

54 metadata: dict[str, Any] = field(default_factory=dict) 

55 

56 

57@dataclass 

58class ExtractedTable: 

59 """Represents an extracted table from a document. 

60 

61 Attributes: 

62 table_id: Unique identifier for the table. 

63 page_number: Page number where table was found. 

64 headers: Column headers. 

65 rows: Table rows (list of lists). 

66 caption: Table caption or title. 

67 format: Format of the table. 

68 metadata: Additional metadata. 

69 """ 

70 

71 table_id: str 

72 page_number: int 

73 headers: list[str] = field(default_factory=list) 

74 rows: list[list[str]] = field(default_factory=list) 

75 caption: str | None = None 

76 format: TableFormat = TableFormat.MARKDOWN 

77 metadata: dict[str, Any] = field(default_factory=dict) 

78 

79 def to_markdown(self) -> str: 

80 """Convert table to Markdown format.""" 

81 if not self.headers and not self.rows: 

82 return "" 

83 

84 lines = [] 

85 

86 # Add caption 

87 if self.caption: 

88 lines.append(f"**{self.caption}**\n") 

89 

90 # Add headers 

91 if self.headers: 

92 lines.append("| " + " | ".join(self.headers) + " |") 

93 lines.append("|" + "|".join(["---"] * len(self.headers)) + "|") 

94 

95 # Add rows 

96 for row in self.rows: 

97 lines.append("| " + " | ".join(row) + " |") 

98 

99 return "\n".join(lines) + "\n" 

100 

101 def to_html(self) -> str: 

102 """Convert table to HTML format.""" 

103 if not self.headers and not self.rows: 

104 return "" 

105 

106 html_parts = [] 

107 

108 # Add caption 

109 if self.caption: 

110 html_parts.append(f"<caption>{self.caption}</caption>") 

111 

112 # Build table 

113 html_parts.append("<table>") 

114 

115 # Add headers 

116 if self.headers: 

117 html_parts.append("<thead><tr>") 

118 for header in self.headers: 

119 html_parts.append(f"<th>{header}</th>") 

120 html_parts.append("</tr></thead>") 

121 

122 # Add rows 

123 html_parts.append("<tbody>") 

124 for row in self.rows: 

125 html_parts.append("<tr>") 

126 for cell in row: 

127 html_parts.append(f"<td>{cell}</td>") 

128 html_parts.append("</tr>") 

129 html_parts.append("</tbody>") 

130 

131 html_parts.append("</table>") 

132 return "".join(html_parts) 

133 

134 def to_csv(self) -> str: 

135 """Convert table to CSV format.""" 

136 if not self.headers and not self.rows: 

137 return "" 

138 

139 lines = [] 

140 

141 # Add headers 

142 if self.headers: 

143 lines.append(",".join(self.headers)) 

144 

145 # Add rows 

146 for row in self.rows: 

147 lines.append(",".join(row)) 

148 

149 return "\n".join(lines) 

150 

151 def to_json(self) -> list[dict[str, Any]]: 

152 """Convert table to JSON-compatible list of dicts.""" 

153 if not self.headers and not self.rows: 

154 return [] 

155 

156 result: list[dict[str, Any]] = [] 

157 if self.headers: 

158 for row in self.rows: 

159 result.append(dict(zip(self.headers, row, strict=False))) 

160 else: 

161 for i, row in enumerate(self.rows): 

162 result.append({"row": i, "data": row}) 

163 

164 return result 

165 

166 

167@dataclass 

168class DocumentMetadata: 

169 """Metadata for a preprocessed document. 

170 

171 Attributes: 

172 doc_type: Type of document. 

173 page_count: Number of pages. 

174 title: Document title. 

175 author: Document author. 

176 created_at: Creation timestamp. 

177 updated_at: Last modification timestamp. 

178 language: Document language. 

179 word_count: Number of words. 

180 char_count: Number of characters. 

181 images: List of extracted images. 

182 tables: List of extracted tables. 

183 custom_metadata: Additional custom metadata. 

184 """ 

185 

186 doc_type: DocumentType = DocumentType.UNKNOWN 

187 page_count: int = 0 

188 title: str | None = None 

189 author: str | None = None 

190 created_at: datetime | None = None 

191 updated_at: datetime | None = None 

192 language: str | None = None 

193 word_count: int = 0 

194 char_count: int = 0 

195 keywords: list[str] = field(default_factory=list) 

196 summary: str | None = None 

197 images: list[ExtractedImage] = field(default_factory=list) 

198 tables: list[ExtractedTable] = field(default_factory=list) 

199 custom_metadata: dict[str, Any] = field(default_factory=dict) 

200 

201 def add_image(self, image: ExtractedImage) -> None: 

202 """Add an extracted image.""" 

203 self.images.append(image) 

204 

205 def add_table(self, table: ExtractedTable) -> None: 

206 """Add an extracted table.""" 

207 self.tables.append(table) 

208 

209 @property 

210 def custom(self) -> dict[str, Any]: 

211 """Backwards compatible alias for custom_metadata.""" 

212 return self.custom_metadata 

213 

214 @property 

215 def document_type(self) -> DocumentType: 

216 """Backwards compatible alias for doc_type.""" 

217 return self.doc_type 

218 

219 def merge(self, other: DocumentMetadata) -> None: 

220 """Merge another metadata object into this one.""" 

221 if other.title and not self.title: 

222 self.title = other.title 

223 if other.author and not self.author: 

224 self.author = other.author 

225 if other.language and not self.language: 

226 self.language = other.language 

227 if other.word_count > 0 and self.word_count == 0: 

228 self.word_count = other.word_count 

229 if other.char_count > 0 and self.char_count == 0: 

230 self.char_count = other.char_count 

231 if ( 

232 other.doc_type != DocumentType.UNKNOWN 

233 and self.doc_type == DocumentType.UNKNOWN 

234 ): 

235 self.doc_type = other.doc_type 

236 

237 self.keywords.extend(k for k in other.keywords if k not in self.keywords) 

238 self.images.extend(other.images) 

239 self.tables.extend(other.tables) 

240 self.custom_metadata.update(other.custom_metadata) 

241 if other.summary and not self.summary: 

242 self.summary = other.summary