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

1"""Loader registry and smart loader for document format auto-detection. 

2 

3Provides ``LoaderRegistry`` (extension/mime-to-loader map) and 

4``SmartLoader`` (auto-detects format and delegates to the right loader). 

5""" 

6 

7from __future__ import annotations 

8 

9from pathlib import Path 

10from typing import Any 

11from urllib.parse import urlparse 

12 

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

14from lexigram.ai.rag.loaders.core import AbstractDocumentLoader 

15from lexigram.ai.rag.types import RAGError 

16from lexigram.logging import ( 

17 get_logger, 

18) 

19from lexigram.primitives.registry import Registry 

20 

21logger = get_logger(__name__) 

22 

23 

24class UnsupportedFormatError(RAGError): 

25 """Raised when no loader is registered for the given file format or URL.""" 

26 

27 _code: str = "LEX_ERR_RAG_016" 

28 

29 def __init__(self, source: str) -> None: 

30 """Create an UnsupportedFormatError. 

31 

32 Args: 

33 source: The source path or URL that could not be matched. 

34 """ 

35 super().__init__(f"No loader registered for source: {source!r}") 

36 self.source = source 

37 

38 

39class LoaderRegistry: 

40 """Registry mapping file extensions and URL schemes to loader instances. 

41 

42 Uses extension strings (including the leading dot) as keys, e.g. 

43 ``.pdf``, ``.json``. The special key ``"url"`` is used for HTTP/HTTPS 

44 sources. 

45 

46 Follows the registry pattern — no ``if/elif`` chains. 

47 

48 Example: 

49 >>> registry = LoaderRegistry() 

50 >>> registry.register([".txt", ".md"], TextLoader()) 

51 >>> registry.register([".json", ".jsonl"], JSONLoader()) 

52 >>> loader = registry.get_loader("report.json") 

53 """ 

54 

55 def __init__(self) -> None: 

56 """Initialize an empty LoaderRegistry.""" 

57 self._registry: Registry[str, AbstractDocumentLoader] = Registry(name="loaders") 

58 

59 def register(self, extensions: list[str], loader: AbstractDocumentLoader) -> None: 

60 """Register a loader for one or more extensions. 

61 

62 Args: 

63 extensions: List of file extensions (e.g. ``[".pdf"]``) or 

64 the special value ``["url"]`` for HTTP/HTTPS sources. 

65 loader: Loader instance to handle those extensions. 

66 """ 

67 for ext in extensions: 

68 self._registry.register(ext.lower(), loader) 

69 

70 def get_loader(self, source: str) -> AbstractDocumentLoader | None: 

71 """Return the loader registered for this source, or None. 

72 

73 Args: 

74 source: File path string or URL. 

75 

76 Returns: 

77 The registered loader, or ``None`` if no match found. 

78 """ 

79 src = str(source) 

80 parsed = urlparse(src) 

81 if parsed.scheme in ("http", "https"): 

82 return self._registry.get("url") 

83 ext = Path(src).suffix.lower() 

84 return self._registry.get(ext) 

85 

86 async def load(self, source: str | Path, **kwargs: Any) -> list[Chunk]: 

87 """Resolve the loader and load the source. 

88 

89 Args: 

90 source: File path or URL. 

91 **kwargs: Passed through to the resolved loader. 

92 

93 Returns: 

94 List of chunks. 

95 

96 Raises: 

97 UnsupportedFormatError: If no loader is registered for this 

98 extension or URL scheme. 

99 RAGError: If the underlying loader raises. 

100 """ 

101 loader = self.get_loader(str(source)) 

102 if loader is None: 

103 raise UnsupportedFormatError(str(source)) 

104 return await loader.load(source, **kwargs) 

105 

106 

107def build_default_registry() -> LoaderRegistry: 

108 """Build a ``LoaderRegistry`` pre-populated with all built-in loaders. 

109 

110 Only loaders whose optional dependencies are available are registered; 

111 a missing dependency causes that loader to be silently skipped. 

112 

113 Returns: 

114 A ready-to-use registry with loaders registered for well-known 

115 extensions. 

116 """ 

117 from lexigram.ai.rag.loaders.core import ( 

118 CSVLoader, 

119 HTMLLoader, 

120 JSONLoader, 

121 MarkdownLoader, 

122 TextLoader, 

123 ) 

124 from lexigram.ai.rag.loaders.p1_loaders import ( 

125 CodeLoader, 

126 DocxLoader, 

127 EmailLoader, 

128 ExcelLoader, 

129 WebScraperLoader, 

130 ) 

131 

132 registry = LoaderRegistry() 

133 

134 # P0 loaders (no optional deps) 

135 registry.register([".txt", ".rst"], TextLoader()) 

136 registry.register([".md", ".markdown"], MarkdownLoader()) 

137 registry.register([".json", ".jsonl"], JSONLoader()) 

138 registry.register([".csv"], CSVLoader()) 

139 registry.register([".tsv"], CSVLoader(delimiter="\t")) 

140 

141 # P0 — beautifulsoup4 optional 

142 try: 

143 registry.register([".html", ".htm"], HTMLLoader()) 

144 registry.register(["url"], WebScraperLoader()) # type: ignore[arg-type] 

145 except ImportError as e: 

146 logger.debug( 

147 "loader_registration_skipped", 

148 loaders=["HTMLLoader", "WebScraperLoader"], 

149 error=str(e), 

150 ) 

151 

152 # P0 — pypdf optional 

153 try: 

154 from lexigram.ai.rag.loaders.core import PDFLoader 

155 

156 registry.register([".pdf"], PDFLoader()) 

157 except ImportError as e: 

158 logger.debug("loader_registration_skipped", loaders=["PDFLoader"], error=str(e)) 

159 

160 # P1 loaders 

161 try: 

162 registry.register([".docx"], DocxLoader()) # type: ignore[arg-type] 

163 except ImportError as e: 

164 logger.debug( 

165 "loader_registration_skipped", loaders=["DocxLoader"], error=str(e) 

166 ) 

167 

168 try: 

169 registry.register([".xlsx", ".xls"], ExcelLoader()) # type: ignore[arg-type] 

170 except ImportError as e: 

171 logger.debug( 

172 "loader_registration_skipped", loaders=["ExcelLoader"], error=str(e) 

173 ) 

174 

175 registry.register([".eml"], EmailLoader()) # type: ignore[arg-type] 

176 

177 code_loader = CodeLoader() 

178 registry.register( 

179 [".py", ".js", ".ts", ".java", ".go", ".rs", ".rb", ".cpp", ".c", ".cs"], 

180 code_loader, # type: ignore[arg-type] 

181 ) 

182 

183 return registry 

184 

185 

186class SmartLoader: 

187 """Auto-detect file format from extension or URL scheme and delegate. 

188 

189 Uses :class:`LoaderRegistry` internally. If no custom registry is 

190 provided, :func:`build_default_registry` is called on first use. 

191 

192 Example: 

193 >>> loader = SmartLoader() 

194 >>> chunks = await loader.load("report.pdf") 

195 >>> chunks = await loader.load("https://example.com/page") 

196 """ 

197 

198 def __init__(self, registry: LoaderRegistry | None = None) -> None: 

199 """Initialize SmartLoader. 

200 

201 Args: 

202 registry: Optional pre-configured registry. Defaults to the 

203 result of :func:`build_default_registry`. 

204 """ 

205 self._registry = registry 

206 

207 @property 

208 def registry(self) -> LoaderRegistry: 

209 """Lazy-initialize the registry on first access.""" 

210 if self._registry is None: 

211 self._registry = build_default_registry() 

212 return self._registry 

213 

214 async def load(self, source: str | Path, **kwargs: Any) -> list[Chunk]: 

215 """Load a source using the auto-detected loader. 

216 

217 Args: 

218 source: File path or URL. 

219 **kwargs: Passed through to the resolved loader. 

220 

221 Returns: 

222 List of chunks. 

223 

224 Raises: 

225 UnsupportedFormatError: If the file extension or URL scheme has 

226 no registered loader. 

227 RAGError: If loading fails. 

228 """ 

229 return await self.registry.load(source, **kwargs)