Coverage for src/local_deep_research/web_search_engines/engines/search_engine_arxiv.py: 97%
162 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1from typing import Any, Dict, List, Optional
3import arxiv
4from langchain_core.language_models import BaseLLM
6from ...constants import SNIPPET_LENGTH_SHORT
7from ...security.secure_logging import logger
8from ..rate_limiting import RateLimitError
9from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity
12class ArXivSearchEngine(BaseSearchEngine):
13 """arXiv search engine implementation with two-phase approach"""
15 # Mark as public search engine
16 is_public = True
17 egress_sensitivity = Sensitivity.NON_SENSITIVE
18 egress_exposure = Exposure.EXPOSING
19 # Not a generic search engine (specialized for academic papers)
20 is_generic = False
21 # Scientific/academic search engine
22 is_scientific = True
23 is_lexical = True
24 needs_llm_relevance_filter = True
26 def __init__(
27 self,
28 max_results: int = 10,
29 sort_by: str = "relevance",
30 sort_order: str = "descending",
31 include_full_text: bool = False,
32 download_dir: Optional[str] = None,
33 max_full_text: int = 1,
34 llm: Optional[BaseLLM] = None,
35 max_filtered_results: Optional[int] = None,
36 settings_snapshot: Optional[Dict[str, Any]] = None,
37 ): # Added this parameter
38 """
39 Initialize the arXiv search engine.
41 Args:
42 max_results: Maximum number of search results
43 sort_by: Sorting criteria ('relevance', 'lastUpdatedDate', or 'submittedDate')
44 sort_order: Sort order ('ascending' or 'descending')
45 include_full_text: Whether to include full paper content in results (downloads PDF)
46 download_dir: Directory to download PDFs to (if include_full_text is True)
47 max_full_text: Maximum number of PDFs to download and process (default: 1)
48 llm: Language model for relevance filtering
49 max_filtered_results: Maximum number of results to keep after filtering
50 settings_snapshot: Settings snapshot for thread context
51 """
52 # Initialize the journal reputation filter if needed.
53 # Runs as a preview filter (before LLM relevance) because Tiers 1-3
54 # are instant data lookups — no point sending irrelevant journals
55 # through the expensive LLM relevance filter.
56 preview_filters = []
57 journal_filter = self._create_journal_filter(
58 "arxiv", llm, settings_snapshot
59 )
60 if journal_filter is not None:
61 preview_filters.append(journal_filter)
63 super().__init__(
64 llm=llm,
65 max_filtered_results=max_filtered_results,
66 max_results=max_results,
67 preview_filters=preview_filters, # type: ignore[arg-type]
68 settings_snapshot=settings_snapshot,
69 )
70 self.max_results = max(self.max_results, 25)
71 self.sort_by = sort_by
72 self.sort_order = sort_order
73 self.include_full_text = include_full_text
74 self.download_dir = download_dir
75 self.max_full_text = max_full_text
77 # Map sort parameters to arxiv package parameters
78 self.sort_criteria = {
79 "relevance": arxiv.SortCriterion.Relevance,
80 "lastUpdatedDate": arxiv.SortCriterion.LastUpdatedDate,
81 "submittedDate": arxiv.SortCriterion.SubmittedDate,
82 }
84 self.sort_directions = {
85 "ascending": arxiv.SortOrder.Ascending,
86 "descending": arxiv.SortOrder.Descending,
87 }
89 def _get_search_results(self, query: str) -> List[Any]:
90 """
91 Helper method to get search results from arXiv API.
93 Args:
94 query: The search query
96 Returns:
97 List of arXiv paper objects
98 """
99 # Configure the search client
100 sort_criteria = self.sort_criteria.get(
101 self.sort_by, arxiv.SortCriterion.Relevance
102 )
103 sort_order = self.sort_directions.get(
104 self.sort_order, arxiv.SortOrder.Descending
105 )
107 # Create the search client
108 client = arxiv.Client(page_size=self.max_results)
110 # Create the search query
111 search = arxiv.Search(
112 query=query,
113 max_results=self.max_results,
114 sort_by=sort_criteria,
115 sort_order=sort_order,
116 )
118 # Apply rate limiting before making the request
119 self._last_wait_time = self.rate_tracker.apply_rate_limit(
120 self.engine_type
121 )
123 # Get the search results
124 return list(client.results(search))
126 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
127 """
128 Get preview information for arXiv papers.
130 Args:
131 query: The search query
133 Returns:
134 List of preview dictionaries
135 """
136 logger.info("Getting paper previews from arXiv")
138 try:
139 # Get search results from arXiv
140 papers = self._get_search_results(query)
142 # Store the paper objects for later use
143 self._papers = {paper.entry_id: paper for paper in papers}
145 # Format results as previews with basic information
146 previews = []
147 for paper in papers:
148 preview = {
149 "id": paper.entry_id, # Use entry_id as ID
150 "title": paper.title,
151 "link": paper.entry_id, # arXiv URL
152 "snippet": (
153 paper.summary[:SNIPPET_LENGTH_SHORT] + "..."
154 if len(paper.summary) > SNIPPET_LENGTH_SHORT
155 else paper.summary
156 ),
157 "authors": [
158 author.name for author in paper.authors[:3]
159 ], # First 3 authors
160 "published": (
161 paper.published.strftime("%Y-%m-%d")
162 if paper.published
163 else None
164 ),
165 "journal_ref": paper.journal_ref,
166 "source": "arXiv",
167 }
169 previews.append(preview)
171 return previews
173 except Exception as e:
174 error_msg = str(e)
175 safe_msg = self._scrub_error(e)
176 logger.exception(
177 f"Error getting arXiv previews ({type(e).__name__}): {safe_msg}"
178 )
180 # Check for rate limiting patterns
181 if (
182 "429" in error_msg
183 or "too many requests" in error_msg.lower()
184 or "rate limit" in error_msg.lower()
185 or "service unavailable" in error_msg.lower()
186 or "503" in error_msg
187 ):
188 # `from None` suppresses the implicit __context__ chain:
189 # the original exception still carries the raw message, so
190 # a full traceback render (chain=True) would re-leak the
191 # secret that safe_msg just scrubbed.
192 raise RateLimitError(
193 f"arXiv rate limit hit: {safe_msg}"
194 ) from None
196 return []
198 def _get_full_content(
199 self, relevant_items: List[Dict[str, Any]]
200 ) -> List[Dict[str, Any]]:
201 """
202 Get full content for the relevant arXiv papers.
203 Downloads PDFs and extracts text when include_full_text is True.
204 Limits the number of PDFs processed to max_full_text.
206 Args:
207 relevant_items: List of relevant preview dictionaries
209 Returns:
210 List of result dictionaries with full content
211 """
212 logger.info("Getting full content for relevant arXiv papers")
214 results = []
215 pdf_count = 0 # Track number of PDFs processed
217 for item in relevant_items:
218 # Start with the preview data
219 result = item.copy()
221 # Get the paper ID
222 paper_id = item.get("id")
224 # Try to get the full paper from our cache
225 paper = None
226 if hasattr(self, "_papers") and paper_id in self._papers:
227 paper = self._papers[paper_id]
229 if paper:
230 # Add complete paper information
231 result.update(
232 {
233 "pdf_url": paper.pdf_url,
234 "authors": [
235 author.name for author in paper.authors
236 ], # All authors
237 "published": (
238 paper.published.strftime("%Y-%m-%d")
239 if paper.published
240 else None
241 ),
242 "updated": (
243 paper.updated.strftime("%Y-%m-%d")
244 if paper.updated
245 else None
246 ),
247 "categories": paper.categories,
248 "summary": paper.summary, # Full summary
249 "comment": paper.comment,
250 "doi": paper.doi,
251 # Explicitly forward for journal quality filter
252 "journal_ref": paper.journal_ref,
253 }
254 )
256 # Default to using summary as content
257 result["content"] = paper.summary
258 result["full_content"] = paper.summary
260 # Download PDF and extract text if requested and within limit
261 if (
262 self.include_full_text
263 and self.download_dir
264 and pdf_count < self.max_full_text
265 ):
266 try:
267 # Download the paper
268 pdf_count += (
269 1 # Increment counter before attempting download
270 )
271 # Apply rate limiting before PDF download
272 self.rate_tracker.apply_rate_limit(self.engine_type)
274 paper_path = paper.download_pdf(
275 dirpath=self.download_dir
276 )
277 result["pdf_path"] = str(paper_path)
279 # Extract text from PDF
280 try:
281 # Try pypdf first
282 try:
283 from pypdf import PdfReader
285 with open(paper_path, "rb") as pdf_file:
286 pdf_reader = PdfReader(pdf_file)
287 pdf_text = ""
288 for page in pdf_reader.pages:
289 pdf_text += page.extract_text() + "\n\n"
291 if (
292 pdf_text.strip()
293 ): # Only use if we got meaningful text
294 result["content"] = pdf_text
295 result["full_content"] = pdf_text
296 logger.info(
297 "Successfully extracted text from PDF using pypdf"
298 )
299 except (ImportError, Exception) as e1:
300 # Fall back to pdfplumber
301 try:
302 import pdfplumber
304 with pdfplumber.open(paper_path) as pdf:
305 pdf_text = ""
306 for plumber_page in pdf.pages:
307 pdf_text += (
308 plumber_page.extract_text()
309 + "\n\n"
310 )
312 if ( 312 ↛ 356line 312 didn't jump to line 356
313 pdf_text.strip()
314 ): # Only use if we got meaningful text
315 result["content"] = pdf_text
316 result["full_content"] = pdf_text
317 logger.info(
318 "Successfully extracted text from PDF using pdfplumber"
319 )
320 except (ImportError, Exception) as e2:
321 safe_e1 = self._scrub_error(e1)
322 safe_e2 = self._scrub_error(e2)
323 logger.exception(
324 f"PDF text extraction failed ({type(e1).__name__}, then {type(e2).__name__}): {safe_e1}, then {safe_e2}"
325 )
326 logger.info(
327 "Using paper summary as content instead"
328 )
329 except Exception as e:
330 safe_msg = self._scrub_error(e)
331 logger.exception(
332 f"Error extracting text from PDF ({type(e).__name__}): {safe_msg}"
333 )
334 logger.info(
335 "Using paper summary as content instead"
336 )
337 except Exception as e:
338 safe_msg = self._scrub_error(e)
339 logger.exception(
340 f"Error downloading paper {paper.title} ({type(e).__name__}): {safe_msg}"
341 )
342 result["pdf_path"] = None
343 pdf_count -= 1 # Decrement counter if download fails
344 elif (
345 self.include_full_text
346 and self.download_dir
347 and pdf_count >= self.max_full_text
348 ):
349 # Reached PDF limit
350 logger.info(
351 f"Maximum number of PDFs ({self.max_full_text}) reached. Skipping remaining PDFs."
352 )
353 result["content"] = paper.summary
354 result["full_content"] = paper.summary
356 results.append(result)
358 return results
360 def run(
361 self, query: str, research_context: Dict[str, Any] | None = None
362 ) -> List[Dict[str, Any]]:
363 """
364 Execute a search using arXiv with the two-phase approach.
366 Args:
367 query: The search query
368 research_context: Context from previous research to use.
370 Returns:
371 List of search results
372 """
373 logger.info("---Execute a search using arXiv---")
375 # Use the implementation from the parent class which handles all phases
376 results = super().run(query, research_context=research_context)
378 # Clean up
379 if hasattr(self, "_papers"):
380 del self._papers
382 return results
384 def get_paper_details(self, arxiv_id: str) -> Dict[str, Any]:
385 """
386 Get detailed information about a specific arXiv paper.
388 Args:
389 arxiv_id: arXiv ID of the paper (e.g., '2101.12345')
391 Returns:
392 Dictionary with paper information
393 """
394 try:
395 # Create the search client
396 client = arxiv.Client()
398 # Search for the specific paper
399 search = arxiv.Search(id_list=[arxiv_id], max_results=1)
401 # Apply rate limiting before fetching paper by ID
402 self._last_wait_time = self.rate_tracker.apply_rate_limit(
403 self.engine_type
404 )
406 # Get the paper
407 papers = list(client.results(search))
408 if not papers:
409 return {}
411 paper = papers[0]
413 # Format result based on config
414 result = {
415 "title": paper.title,
416 "link": paper.entry_id,
417 "snippet": (
418 paper.summary[:250] + "..."
419 if len(paper.summary) > 250
420 else paper.summary
421 ),
422 "authors": [
423 author.name for author in paper.authors[:3]
424 ], # First 3 authors
425 "journal_ref": paper.journal_ref,
426 }
428 result.update(
429 {
430 "pdf_url": paper.pdf_url,
431 "authors": [
432 author.name for author in paper.authors
433 ], # All authors
434 "published": (
435 paper.published.strftime("%Y-%m-%d")
436 if paper.published
437 else None
438 ),
439 "updated": (
440 paper.updated.strftime("%Y-%m-%d")
441 if paper.updated
442 else None
443 ),
444 "categories": paper.categories,
445 "summary": paper.summary, # Full summary
446 "comment": paper.comment,
447 "doi": paper.doi,
448 "content": paper.summary, # Use summary as content
449 "full_content": paper.summary, # For consistency
450 }
451 )
453 # Download PDF if requested
454 if self.include_full_text and self.download_dir:
455 try:
456 # Apply rate limiting before PDF download
457 self.rate_tracker.apply_rate_limit(self.engine_type)
459 # Download the paper
460 paper_path = paper.download_pdf(dirpath=self.download_dir)
461 result["pdf_path"] = str(paper_path)
462 except Exception as e:
463 safe_msg = self._scrub_error(e)
464 logger.exception(
465 f"Error downloading paper ({type(e).__name__}): {safe_msg}"
466 )
468 return result
470 except Exception as e:
471 safe_msg = self._scrub_error(e)
472 logger.exception(
473 f"Error getting paper details ({type(e).__name__}): {safe_msg}"
474 )
475 return {}
477 def search_by_author(
478 self, author_name: str, max_results: Optional[int] = None
479 ) -> List[Dict[str, Any]]:
480 """
481 Search for papers by a specific author.
483 Args:
484 author_name: Name of the author
485 max_results: Maximum number of results (defaults to self.max_results)
487 Returns:
488 List of papers by the author
489 """
490 original_max_results = self.max_results
492 try:
493 if max_results:
494 self.max_results = max_results
496 query = f'au:"{author_name}"'
497 return self.run(query)
499 finally:
500 # Restore original value
501 self.max_results = original_max_results
503 def search_by_category(
504 self, category: str, max_results: Optional[int] = None
505 ) -> List[Dict[str, Any]]:
506 """
507 Search for papers in a specific arXiv category.
509 Args:
510 category: arXiv category (e.g., 'cs.AI', 'physics.optics')
511 max_results: Maximum number of results (defaults to self.max_results)
513 Returns:
514 List of papers in the category
515 """
516 original_max_results = self.max_results
518 try:
519 if max_results:
520 self.max_results = max_results
522 query = f"cat:{category}"
523 return self.run(query)
525 finally:
526 # Restore original value
527 self.max_results = original_max_results