Coverage for src/local_deep_research/utilities/arxiv_api.py: 90%

1"""Shared arXiv API transport policy (issue #4783). 

2 

3Sole owner of the arXiv API request policy: every legacy arXiv API query 

4request must traverse one process-wide, monotonic, at least three-second 

5single-flight gate and travel through a :class:`SafeSession` with a pinned 

6timeout and redirects disabled. 

7""" 

8 

9from __future__ import annotations 

10 

11import threading 

12import time 

13from collections.abc import Callable, Iterable, Sequence 

14from contextlib import contextmanager 

15from dataclasses import dataclass 

16from datetime import datetime 

17from enum import StrEnum 

18from typing import ( 

19 TYPE_CHECKING, 

20 Final, 

21 Protocol, 

22 TypeAlias, 

23 assert_never, 

24 override, 

25) 

26 

27import arxiv 

28import requests 

29from requests import PreparedRequest 

30 

31from ..security import SafeSession 

32from .resource_utils import safe_close 

33 

34if TYPE_CHECKING: 

35 from collections.abc import Generator 

36 

37 from requests import Response 

38 

39# arXiv API Terms of Use: one connection at a time, with at least three 

40# seconds between request starts 

41# (https://info.arxiv.org/help/api/tou.html). This gate is intentionally 

42# process-local: it serializes request starts within a single LDR process 

43# only and provides no machine-wide enforcement. Deployments running 

44# multiple LDR processes must coordinate those processes externally. 

45_ARXIV_API_REQUEST_INTERVAL_SECONDS: Final = 3.0 

46_arxiv_api_request_lock = threading.Lock() 

47_last_request_started_at: float | None = None 

48_monotonic = time.monotonic 

49_sleep = time.sleep 

50 

51_ARXIV_API_TIMEOUT_SECONDS: Final = 10 

52 

53 

54class ArxivSortCriterion(StrEnum): 

55 RELEVANCE = "relevance" 

56 LAST_UPDATED_DATE = "lastUpdatedDate" 

57 SUBMITTED_DATE = "submittedDate" 

58 

59 

60class ArxivSortOrder(StrEnum): 

61 ASCENDING = "ascending" 

62 DESCENDING = "descending" 

63 

64 

65@dataclass(frozen=True, slots=True) 

66class ArxivQueryRequest: 

67 query: str 

68 max_results: int 

69 sort_by: ArxivSortCriterion = ArxivSortCriterion.RELEVANCE 

70 sort_order: ArxivSortOrder = ArxivSortOrder.DESCENDING 

71 

72 

73@dataclass(frozen=True, slots=True) 

74class ArxivIdRequest: 

75 arxiv_id: str 

76 

77 

78type ArxivRequest = ArxivQueryRequest | ArxivIdRequest 

79 

80 

81class ArxivAuthor(Protocol): 

82 name: str 

83 

84 

85class ArxivPaper(Protocol): 

86 entry_id: str 

87 title: str 

88 summary: str 

89 published: datetime 

90 updated: datetime 

91 categories: list[str] 

92 comment: str | None 

93 doi: str | None 

94 journal_ref: str | None 

95 

96 @property 

97 def authors(self) -> Sequence[ArxivAuthor]: ... 97 ↛ anywhereline 97 didn't jump anywhere: it always raised an exception.

98 

99 @property 

100 def pdf_url(self) -> str | None: ... 100 ↛ exitline 100 didn't return from function 'pdf_url' because

101 

102 

103_SORT_CRITERIA: Final = { 

104 ArxivSortCriterion.RELEVANCE: arxiv.SortCriterion.Relevance, 

105 ArxivSortCriterion.LAST_UPDATED_DATE: arxiv.SortCriterion.LastUpdatedDate, 

106 ArxivSortCriterion.SUBMITTED_DATE: arxiv.SortCriterion.SubmittedDate, 

107} 

108_SORT_ORDERS: Final = { 

109 ArxivSortOrder.ASCENDING: arxiv.SortOrder.Ascending, 

110 ArxivSortOrder.DESCENDING: arxiv.SortOrder.Descending, 

111} 

112 

113 

114class _Readable(Protocol): 

115 def read(self, length: int = ..., /) -> str | bytes: ... 115 ↛ exitline 115 didn't return from function 'read' because

116 

117 

118class _ItemSource(Protocol): 

119 def items(self) -> Iterable[tuple[_GetArgument, _GetArgument]]: ... 119 ↛ exitline 119 didn't return from function 'items' because

120 

121 

122_GetArgument: TypeAlias = ( 

123 None 

124 | bool 

125 | int 

126 | float 

127 | str 

128 | bytes 

129 | PreparedRequest 

130 | Iterable["_GetArgument"] 

131 | _Readable 

132 | _ItemSource 

133 | Callable[..., "_GetArgument"] 

134) 

135 

136 

137@contextmanager 

138def arxiv_api_request_gate() -> Generator[None, None, None]: 

139 """NOT REENTRANT -- never enter this from code that may already 

140 be running under the gate. The lock is a plain threading.Lock held 

141 across both the sleep and the request body, so a nested entry 

142 deadlocks silently rather than raising. _GatedArxivSession.get() 

143 enters it, so anything handed that session as its transport must 

144 not wrap its own call in the gate. 

145 

146 Run one arXiv API request start under the process-wide policy gate. 

147 

148 Entering blocks until at least ``_ARXIV_API_REQUEST_INTERVAL_SECONDS`` 

149 monotonic seconds have passed since the previous start, then stamps the 

150 new start while holding the single-flight lock across the whole body. 

151 Failed starts still count for spacing, and the lock is always released. 

152 """ 

153 global _last_request_started_at 

154 with _arxiv_api_request_lock: 

155 while _last_request_started_at is not None: 

156 remaining = ( 

157 _last_request_started_at 

158 + _ARXIV_API_REQUEST_INTERVAL_SECONDS 

159 - _monotonic() 

160 ) 

161 if remaining <= 0: 

162 break 

163 _sleep(remaining) 

164 _last_request_started_at = _monotonic() 

165 yield 

166 

167 

168class _GatedArxivSession(SafeSession): 

169 """SafeSession routing arXiv API wire calls through the shared gate.""" 

170 

171 @override 

172 def get( 

173 self, 

174 url: str | bytes, 

175 params: _GetArgument = None, 

176 **kwargs: _GetArgument, 

177 ) -> Response: 

178 request_kwargs: dict[str, _GetArgument] = { 

179 **kwargs, 

180 "timeout": _ARXIV_API_TIMEOUT_SECONDS, 

181 "allow_redirects": False, 

182 } 

183 match url: 

184 case str(): 184 ↛ 186line 184 didn't jump to line 186 because the pattern on line 184 always matched

185 request_url = url 

186 case bytes(): 

187 request_url = url.decode() 

188 case unreachable: 

189 assert_never(unreachable) 

190 with arxiv_api_request_gate(): 

191 return super().request( 

192 "GET", request_url, params=params, **request_kwargs 

193 ) 

194 

195 

196@contextmanager 

197def _policy_arxiv_client( 

198 *, page_size: int | None = None 

199) -> Generator[arxiv.Client, None, None]: 

200 """Yield an ``arxiv.Client`` adapted to the shared arXiv API policy. 

201 

202 The client is constructed through the ``arxiv.Client`` module attribute 

203 with ``delay_seconds=0`` because the shared gate owns spacing. 

204 :func:`_install_gated_session` then closes the discarded default 

205 session and pins a gated :class:`_GatedArxivSession` onto the 

206 client, raising ``ImportError`` if the client's private ``_session`` 

207 attribute is missing or not a ``requests.Session``. The adapted 

208 session is closed on exit for both normal and exceptional scopes. 

209 """ 

210 constructor_kwargs: dict[str, int] = {"delay_seconds": 0} 

211 if page_size is not None: 

212 constructor_kwargs["page_size"] = page_size 

213 client = arxiv.Client(**constructor_kwargs) 

214 adapted_session = _install_gated_session(client) 

215 try: 

216 yield client 

217 finally: 

218 safe_close(adapted_session, "adapted arXiv client session") 

219 

220 

221def _install_gated_session(client: arxiv.Client) -> _GatedArxivSession: 

222 """Install the gate-enforcing session on ``client`` or fail loudly. 

223 

224 ``arxiv.Client`` keeps its HTTP transport on the private ``_session`` 

225 attribute (a ``requests.Session`` in arxiv 2.4.0, 2.4.1 and 4.0.1; 

226 pyproject.toml does not pin the version for this guard's sake, since 

227 the guard fails closed). A future release renaming or retyping that 

228 attribute would leave the client's default session in place and 

229 silently bypass the shared rate-limit gate, so the mutation is 

230 guarded: anything other than a ``requests.Session`` at ``_session`` 

231 raises ``ImportError`` with maintainer guidance instead of shipping 

232 an ungated client. The discarded default session is closed, a 

233 :class:`_GatedArxivSession` is installed in its place, and it is 

234 returned so the caller can close it on scope exit. 

235 """ 

236 discarded_session = getattr(client, "_session", None) 

237 if not isinstance(discarded_session, requests.Session): 

238 raise ImportError( 

239 "arxiv.Client no longer exposes a requests.Session at " 

240 "_session, so the shared arXiv rate-limit gate cannot be " 

241 "installed. The installed arxiv release changed the client " 

242 "internals; adapt utilities/arxiv_api.py to them, or constrain " 

243 "arxiv to a release that keeps a requests.Session at _session " 

244 "(2.4.0, 2.4.1 and 4.0.1 do).", 

245 ) 

246 safe_close(discarded_session, "discarded arXiv client session") 

247 adapted_session = _GatedArxivSession() 

248 client._session = adapted_session 

249 return adapted_session 

250 

251 

252def _strip_optional(text: str | None) -> str | None: 

253 return None if text is None else text.strip() 

254 

255 

256def _strip_feed_text(paper: arxiv.Result) -> arxiv.Result: 

257 """Strip surrounding whitespace from the text fields LDR reads. 

258 

259 arxiv 2.x parsed the Atom feed with feedparser, which strips the 

260 leading and trailing whitespace of each of these fields. arxiv 4.x 

261 parses it with its own lxml parser and keeps the raw element text, 

262 and the arXiv API returns some of it with surrounding whitespace 

263 (some abstracts start with two spaces). The id, title, abstract, 

264 comment, journal reference, DOI and author names are stripped in 

265 place with ``str.strip()``, as feedparser stripped them, so snippets, 

266 content and author lists carry none of the whitespace 2.x removed. 

267 The SDK already collapses each whitespace run in the title to one 

268 space, which leaves at most one space at each end to strip. 

269 """ 

270 paper.entry_id = paper.entry_id.strip() 

271 paper.title = paper.title.strip() 

272 paper.summary = paper.summary.strip() 

273 paper.comment = _strip_optional(paper.comment) 

274 paper.journal_ref = _strip_optional(paper.journal_ref) 

275 paper.doi = _strip_optional(paper.doi) 

276 for author in paper.authors: 

277 author.name = author.name.strip() 

278 return paper 

279 

280 

281def fetch_arxiv_results(request: ArxivRequest) -> list[ArxivPaper]: 

282 """Fetch and materialize one typed arXiv request under shared policy. 

283 

284 The text fields LDR reads from each result are stripped of 

285 surrounding whitespace (``_strip_feed_text``). 

286 """ 

287 match request: 

288 case ArxivQueryRequest(): 

289 search = arxiv.Search( 

290 query=request.query, 

291 max_results=request.max_results, 

292 sort_by=_SORT_CRITERIA[request.sort_by], 

293 sort_order=_SORT_ORDERS[request.sort_order], 

294 ) 

295 with _policy_arxiv_client(page_size=request.max_results) as client: 

296 return list[ArxivPaper]( 

297 map(_strip_feed_text, client.results(search)) 

298 ) 

299 case ArxivIdRequest(): 299 ↛ 305line 299 didn't jump to line 305 because the pattern on line 299 always matched

300 search = arxiv.Search(id_list=[request.arxiv_id], max_results=1) 

301 with _policy_arxiv_client() as client: 

302 return list[ArxivPaper]( 

303 map(_strip_feed_text, client.results(search)) 

304 ) 

305 case unreachable: 

306 assert_never(unreachable)