Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/loaders/p1_loaders.py: 21%

1"""P1 document loaders for RAG. 

2 

3Compatibility facade for office, web, code, and SQL loaders.""" 

4 

5from __future__ import annotations 

6 

7from pathlib import Path 

8import re 

9from typing import Any 

10 

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

12from lexigram.ai.rag.loaders._io_utils import read_file_text 

13from lexigram.ai.rag.loaders.office import DocxLoader, EmailLoader, ExcelLoader 

14from lexigram.ai.rag.loaders.web import WebScraperLoader 

15from lexigram.ai.rag.types import RAGError 

16from lexigram.contracts.data import DatabaseProviderProtocol 

17 

18# --------------------------------------------------------------------------- 

19# CodeLoader 

20# --------------------------------------------------------------------------- 

21 

22# Regex patterns that detect the start of a top-level definition per language 

23_CODE_SPLIT_PATTERNS: dict[str, str] = { 

24 ".py": r"^(?:async\s+)?(?:def|class)\s+\w+", 

25 ".js": r"^(?:function\s+\w+|class\s+\w+|const\s+\w+\s*=\s*(?:async\s+)?(?:function|\(|[a-zA-Z_]))", 

26 ".ts": r"^(?:export\s+)?(?:async\s+)?(?:function\s+\w+|class\s+\w+|const\s+\w+\s*=)", 

27 ".java": r"^(?:public|private|protected|static|\s)+(?:class|interface|enum|\w+\s*\()", 

28 ".go": r"^func\s+", 

29 ".rs": r"^(?:pub\s+)?(?:fn|struct|enum|impl|trait|mod)\s+\w+", 

30 ".rb": r"^(?:def|class|module)\s+\w+", 

31} 

32 

33 

34class CodeLoader: 

35 """Load source code files with language-aware chunking. 

36 

37 Splits at top-level function/class definitions using per-language 

38 regex patterns. Falls back to fixed-size chunking for unknown languages. 

39 

40 No external dependencies required (optional: tree-sitter for precise AST 

41 parsing in future versions). 

42 """ 

43 

44 def __init__(self, *, max_chunk_lines: int = 100) -> None: 

45 """Initialize code loader. 

46 

47 Args: 

48 max_chunk_lines: Maximum number of lines per chunk for fallback 

49 fixed-size splitting (unknown language extension). 

50 """ 

51 self.max_chunk_lines = max_chunk_lines 

52 

53 async def load(self, source: str | Path) -> list[Chunk]: 

54 """Load a source code file. 

55 

56 Args: 

57 source: Path to the source code file. 

58 

59 Returns: 

60 List of chunks — one per top-level definition or fixed-size block. 

61 

62 Raises: 

63 RAGError: If the file cannot be read. 

64 """ 

65 try: 

66 path = Path(source) 

67 content = await read_file_text(path, encoding="utf-8") 

68 ext = path.suffix.lower() 

69 pattern = _CODE_SPLIT_PATTERNS.get(ext) 

70 

71 chunks: list[Chunk] = [] 

72 

73 if pattern: 

74 split_re = re.compile(pattern, re.MULTILINE) 

75 lines = content.splitlines(keepends=True) 

76 split_indices = [0] 

77 for match in split_re.finditer(content): 

78 # find line number of the match 

79 line_no = content[: match.start()].count("\n") 

80 if line_no > 0 and line_no not in split_indices: 

81 split_indices.append(line_no) 

82 split_indices.append(len(lines)) 

83 

84 for i in range(len(split_indices) - 1): 

85 block = "".join(lines[split_indices[i] : split_indices[i + 1]]) 

86 if block.strip(): 

87 chunks.append( 

88 Chunk( 

89 text=block, 

90 source=str(path), 

91 chunk_index=len(chunks), 

92 metadata={ 

93 "source": str(path), 

94 "type": "code", 

95 "language": ext.lstrip("."), 

96 "start_line": split_indices[i] + 1, 

97 }, 

98 ) 

99 ) 

100 else: 

101 # Fallback: fixed-size line batches 

102 lines = content.splitlines(keepends=True) 

103 for batch_start in range(0, len(lines), self.max_chunk_lines): 

104 block = "".join( 

105 lines[batch_start : batch_start + self.max_chunk_lines] 

106 ) 

107 if block.strip(): 

108 chunks.append( 

109 Chunk( 

110 text=block, 

111 source=str(path), 

112 chunk_index=len(chunks), 

113 metadata={ 

114 "source": str(path), 

115 "type": "code", 

116 "language": ext.lstrip("."), 

117 "start_line": batch_start + 1, 

118 }, 

119 ) 

120 ) 

121 

122 return chunks 

123 

124 except (OSError, UnicodeDecodeError) as e: 

125 msg = f"Failed to read code file {source}: {e}" 

126 raise RAGError(msg) from e 

127 except Exception as e: 

128 msg = f"Unexpected error loading code file {source}: {e}" 

129 raise RAGError(msg) from e 

130 

131 

132# --------------------------------------------------------------------------- 

133# SQLLoader 

134# --------------------------------------------------------------------------- 

135 

136 

137class SQLLoader: 

138 """Load data from a SQL query result as document chunks. 

139 

140 Uses ``DatabaseProviderProtocol`` for all database access — no direct 

141 driver imports. Each row becomes a chunk (or rows are batched). 

142 

143 The caller is responsible for registering the provider and injecting it 

144 via constructor injection. 

145 """ 

146 

147 def __init__( 

148 self, 

149 db: DatabaseProviderProtocol, 

150 *, 

151 query: str, 

152 params: dict[str, Any] | None = None, 

153 text_column: str | None = None, 

154 batch_size: int = 1, 

155 table_name: str = "query", 

156 ) -> None: 

157 """Initialize SQL loader. 

158 

159 Args: 

160 db: Database provider (injected via DI container). 

161 query: SQL SELECT statement to execute. 

162 params: Optional query parameters. 

163 text_column: Column whose value is used as chunk text. When not 

164 set, all columns are joined into a key: value string. 

165 batch_size: Number of rows per chunk. 

166 table_name: Label used in metadata ``source`` field. 

167 """ 

168 self._db = db 

169 self._query = query 

170 self._params = params or {} 

171 self._text_column = text_column 

172 self._batch_size = batch_size 

173 self._table_name = table_name 

174 

175 async def load(self, source: str | Path = "") -> list[Chunk]: 

176 """Execute the SQL query and return a chunk per row (or batch). 

177 

178 Args: 

179 source: Ignored — the data source is the DB connection passed at 

180 construction. Kept for interface compatibility. 

181 

182 Returns: 

183 List of chunks, one per row or per batch of rows. 

184 

185 Raises: 

186 RAGError: If the query fails. 

187 """ 

188 try: 

189 async with self._db.scoped_context(): 

190 conn = await self._db.get_scoped_connection() 

191 rows: list[dict[str, Any]] = await conn.fetch( 

192 self._query, **self._params 

193 ) 

194 

195 label = f"sql://{self._table_name}" 

196 chunks: list[Chunk] = [] 

197 batch: list[dict[str, Any]] = [] 

198 

199 def _row_text(row: dict[str, Any]) -> str: 

200 if self._text_column: 

201 return str(row.get(self._text_column, "")) 

202 return " ".join(f"{k}: {v}" for k, v in row.items()) 

203 

204 for row_idx, row in enumerate(rows): 

205 batch.append(dict(row)) 

206 if len(batch) >= self._batch_size: 

207 text = "\n".join(_row_text(r) for r in batch) 

208 chunk_idx = row_idx // self._batch_size 

209 chunks.append( 

210 Chunk( 

211 text=text, 

212 source=label, 

213 chunk_index=chunk_idx, 

214 metadata={ 

215 "source": label, 

216 "type": "sql", 

217 "table": self._table_name, 

218 "row_start": chunk_idx * self._batch_size, 

219 "row_end": row_idx, 

220 }, 

221 ) 

222 ) 

223 batch = [] 

224 

225 if batch: 

226 chunk_idx = len(rows) // self._batch_size 

227 text = "\n".join(_row_text(r) for r in batch) 

228 chunks.append( 

229 Chunk( 

230 text=text, 

231 source=label, 

232 chunk_index=chunk_idx, 

233 metadata={ 

234 "source": label, 

235 "type": "sql", 

236 "table": self._table_name, 

237 "row_start": chunk_idx * self._batch_size, 

238 "row_end": len(rows) - 1, 

239 }, 

240 ) 

241 ) 

242 

243 return chunks 

244 

245 except RAGError: 

246 raise 

247 except Exception as e: 

248 msg = f"SQL query failed for '{self._table_name}': {e}" 

249 raise RAGError(msg) from e 

250 

251 

252__all__ = [ 

253 "CodeLoader", 

254 "DocxLoader", 

255 "EmailLoader", 

256 "ExcelLoader", 

257 "SQLLoader", 

258 "WebScraperLoader", 

259]