| 1 | """Centralized document fetching for the document_query plugin.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import asyncio |
| 6 | import mimetypes |
| 7 | import os |
| 8 | import tempfile |
| 9 | from contextlib import contextmanager |
| 10 | from dataclasses import dataclass |
| 11 | from pathlib import Path |
| 12 | from typing import Awaitable, Callable |
| 13 | from urllib.parse import urlparse |
| 14 | |
| 15 | from helpers import files |
| 16 | from helpers.network import fetch_public_http_resource |
| 17 | |
| 18 | |
| 19 | InterventionCallback = Callable[[], Awaitable[None]] |
| 20 | |
| 21 | |
| 22 | @dataclass(frozen=True) |
| 23 | class FetchedDocument: |
| 24 | """Fetched document bytes plus metadata needed by parsers.""" |
| 25 | |
| 26 | uri: str |
| 27 | scheme: str |
| 28 | mimetype: str |
| 29 | content: bytes |
| 30 | encoding: str | None = None |
| 31 | charset: str | None = None |
| 32 | local_path: str | None = None |
| 33 | source_uri: str | None = None |
| 34 | |
| 35 | def text(self) -> str: |
| 36 | charset = self.charset or "utf-8" |
| 37 | return self.content.decode(charset, errors="replace") |
| 38 | |
| 39 | def suffix(self) -> str: |
| 40 | path = self.local_path or urlparse(self.uri).path or self.uri |
| 41 | suffix = Path(path).suffix |
| 42 | if suffix: |
| 43 | return suffix |
| 44 | guessed = mimetypes.guess_extension(self.mimetype) |
| 45 | return guessed or ".bin" |
| 46 | |
| 47 | @contextmanager |
| 48 | def local_file(self): |
| 49 | """Yield a filesystem path for parsers that cannot consume bytes.""" |
| 50 | if self.local_path and os.path.exists(self.local_path): |
| 51 | yield self.local_path |
| 52 | return |
| 53 | |
| 54 | tmp = "" |
| 55 | try: |
| 56 | with tempfile.NamedTemporaryFile(delete=False, suffix=self.suffix()) as f: |
| 57 | f.write(self.content) |
| 58 | tmp = f.name |
| 59 | yield tmp |
| 60 | finally: |
| 61 | if tmp and os.path.exists(tmp): |
| 62 | os.unlink(tmp) |
| 63 | |
| 64 | |
| 65 | ProtocolHandler = Callable[ |
| 66 | [str, str, dict, InterventionCallback | None], Awaitable[FetchedDocument] |
| 67 | ] |
| 68 | |
| 69 | _PROTOCOL_HANDLERS: dict[str, ProtocolHandler] = {} |
| 70 | |
| 71 | |
| 72 | def register_protocol_handler(scheme: str, handler: ProtocolHandler) -> None: |
| 73 | """Register or replace a fetch handler for a URI scheme.""" |
| 74 | _PROTOCOL_HANDLERS[scheme.lower()] = handler |
| 75 | |
| 76 | |
| 77 | async def fetch_public_resource( |
| 78 | uri: str, |
| 79 | config: dict | None = None, |
| 80 | intervention_callback: InterventionCallback | None = None, |
| 81 | ) -> FetchedDocument: |
| 82 | """Fetch local or remote content once, then pass bytes to parsers.""" |
| 83 | config = config or {} |
| 84 | parsed = urlparse(uri) |
| 85 | scheme = (parsed.scheme or "file").lower() |
| 86 | handler = _PROTOCOL_HANDLERS.get(scheme) |
| 87 | if not handler: |
| 88 | raise ValueError(f"Unsupported document scheme: {scheme}") |
| 89 | return await handler(uri, scheme, config, intervention_callback) |
| 90 | |
| 91 | |
| 92 | async def _fetch_file( |
| 93 | uri: str, |
| 94 | scheme: str, |
| 95 | config: dict, |
| 96 | intervention_callback: InterventionCallback | None, |
| 97 | ) -> FetchedDocument: |
| 98 | parsed = urlparse(uri) |
| 99 | raw_path = parsed.path if parsed.scheme == "file" else uri |
| 100 | if not raw_path: |
| 101 | raise ValueError(f"Invalid document path: {uri}") |
| 102 | |
| 103 | path = _fix_file_path(raw_path) |
| 104 | mimetype, encoding = mimetypes.guess_type(path) |
| 105 | if encoding: |
| 106 | raise ValueError(f"Compressed documents are unsupported '{encoding}' ({uri})") |
| 107 | mimetype = mimetype or "application/octet-stream" |
| 108 | if mimetype == "application/octet-stream": |
| 109 | raise ValueError(f"Unsupported document mimetype '{mimetype}' ({uri})") |
| 110 | |
| 111 | if intervention_callback: |
| 112 | await intervention_callback() |
| 113 | return FetchedDocument( |
| 114 | uri=path, |
| 115 | source_uri=uri, |
| 116 | scheme=scheme, |
| 117 | mimetype=mimetype, |
| 118 | encoding=encoding, |
| 119 | content=files.read_file_bin(path), |
| 120 | local_path=path, |
| 121 | ) |
| 122 | |
| 123 | |
| 124 | async def _fetch_http( |
| 125 | uri: str, |
| 126 | scheme: str, |
| 127 | config: dict, |
| 128 | intervention_callback: InterventionCallback | None, |
| 129 | ) -> FetchedDocument: |
| 130 | timeout = float(config.get("fetch_timeout", 30)) |
| 131 | retries = max(1, int(config.get("fetch_retries", 3))) |
| 132 | retry_backoff = float(config.get("fetch_retry_backoff", 1.0)) |
| 133 | max_remote_bytes = int(config.get("max_remote_bytes", 50 * 1024 * 1024)) |
| 134 | parsed = urlparse(uri) |
| 135 | guessed_mimetype, encoding = mimetypes.guess_type(parsed.path or uri) |
| 136 | if encoding: |
| 137 | raise ValueError(f"Compressed documents are unsupported '{encoding}' ({uri})") |
| 138 | |
| 139 | last_error = "" |
| 140 | for attempt in range(retries): |
| 141 | try: |
| 142 | if intervention_callback: |
| 143 | await intervention_callback() |
| 144 | resource = await asyncio.to_thread( |
| 145 | fetch_public_http_resource, |
| 146 | uri, |
| 147 | max_bytes=max_remote_bytes, |
| 148 | timeout=(timeout, timeout), |
| 149 | ) |
| 150 | if intervention_callback: |
| 151 | await intervention_callback() |
| 152 | |
| 153 | mimetype = ( |
| 154 | resource.content_type |
| 155 | or guessed_mimetype |
| 156 | or "application/octet-stream" |
| 157 | ) |
| 158 | if mimetype == "application/octet-stream": |
| 159 | raise ValueError( |
| 160 | f"Unsupported document mimetype '{mimetype}' ({uri})" |
| 161 | ) |
| 162 | |
| 163 | return FetchedDocument( |
| 164 | uri=resource.url, |
| 165 | source_uri=uri, |
| 166 | scheme=urlparse(resource.url).scheme or scheme, |
| 167 | mimetype=mimetype, |
| 168 | encoding=encoding, |
| 169 | charset=resource.encoding, |
| 170 | content=resource.content, |
| 171 | ) |
| 172 | except Exception as e: |
| 173 | last_error = str(e) |
| 174 | if attempt < retries - 1: |
| 175 | await asyncio.sleep(retry_backoff) |
| 176 | if intervention_callback: |
| 177 | await intervention_callback() |
| 178 | |
| 179 | raise ValueError(f"Document fetch error: {uri} ({last_error})") |
| 180 | |
| 181 | register_protocol_handler("file", _fetch_file) |
| 182 | register_protocol_handler("http", _fetch_http) |
| 183 | register_protocol_handler("https", _fetch_http) |
| 184 | |
| 185 | |
| 186 | def _fix_file_path(path: str) -> str: |
| 187 | if os.path.isabs(path) and os.path.exists(path): |
| 188 | return path |
| 189 | return files.fix_dev_path(path) |