Coverage for src/local_deep_research/web_search_engines/engines/search_engine_pubchem.py: 94%

261 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-22 21:20 +0000

1"""PubChem search engine for chemical compound information.""" 

2 

3from typing import Any, Dict, List, Optional 

4from urllib.parse import quote 

5 

6from langchain_core.language_models import BaseLLM 

7 

8from ...constants import USER_AGENT 

9from ...security.safe_requests import safe_get 

10from ...security.secure_logging import logger 

11from ..rate_limiting import RateLimitError 

12from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity 

13 

14 

15class PubChemSearchEngine(BaseSearchEngine): 

16 """ 

17 PubChem search engine for chemical compound information. 

18 

19 Provides access to chemical structures, properties, and bioactivity data. 

20 No authentication required. 

21 """ 

22 

23 is_public = True 

24 egress_sensitivity = Sensitivity.NON_SENSITIVE 

25 egress_exposure = Exposure.EXPOSING 

26 is_generic = False 

27 is_scientific = True 

28 is_code = False 

29 is_lexical = True 

30 needs_llm_relevance_filter = True 

31 

32 def __init__( 

33 self, 

34 max_results: int = 10, 

35 include_synonyms: bool = True, 

36 llm: Optional[BaseLLM] = None, 

37 max_filtered_results: Optional[int] = None, 

38 settings_snapshot: Optional[Dict[str, Any]] = None, 

39 **kwargs, 

40 ): 

41 """ 

42 Initialize the PubChem search engine. 

43 

44 Args: 

45 max_results: Maximum number of search results 

46 include_synonyms: Whether to include compound synonyms 

47 llm: Language model for relevance filtering 

48 max_filtered_results: Maximum results after filtering 

49 settings_snapshot: Settings snapshot for thread context 

50 """ 

51 super().__init__( 

52 llm=llm, 

53 max_filtered_results=max_filtered_results, 

54 max_results=max_results, 

55 settings_snapshot=settings_snapshot, 

56 **kwargs, 

57 ) 

58 

59 self.include_synonyms = include_synonyms 

60 self.base_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug" 

61 self.autocomplete_url = ( 

62 "https://pubchem.ncbi.nlm.nih.gov/rest/autocomplete" 

63 ) 

64 

65 # User-Agent header for API requests 

66 self.headers = {"User-Agent": USER_AGENT} 

67 

68 def _search_compounds(self, query: str) -> List[str]: 

69 """Search for compound names matching the query.""" 

70 try: 

71 url = ( 

72 f"{self.autocomplete_url}/compound/{quote(query, safe='')}/json" 

73 ) 

74 params = {"limit": self.max_results * 2} # Get extra for filtering 

75 

76 response = safe_get( 

77 url, params=params, headers=self.headers, timeout=30 

78 ) 

79 self._raise_if_rate_limit(response.status_code) 

80 response.raise_for_status() 

81 data = response.json() 

82 

83 terms: list[str] = data.get("dictionary_terms", {}).get( 

84 "compound", [] 

85 ) 

86 return terms 

87 

88 except RateLimitError: 

89 raise 

90 except Exception as e: 

91 safe_msg = self._scrub_error(e) 

92 logger.exception( 

93 f"PubChem autocomplete search failed ({type(e).__name__}): {safe_msg}" 

94 ) 

95 return [] 

96 

97 def _get_compound_by_name(self, name: str) -> Optional[Dict[str, Any]]: 

98 """Get compound information by name.""" 

99 try: 

100 self.rate_tracker.apply_rate_limit(self.engine_type) 

101 # Get CID first 

102 url = f"{self.base_url}/compound/name/{quote(name, safe='')}/cids/JSON" 

103 response = safe_get(url, headers=self.headers, timeout=30) 

104 

105 if response.status_code == 404: 

106 return None 

107 self._raise_if_rate_limit(response.status_code) 

108 

109 response.raise_for_status() 

110 data = response.json() 

111 cids = data.get("IdentifierList", {}).get("CID", []) 

112 

113 if not cids: 

114 return None 

115 

116 cid = cids[0] 

117 

118 # Get compound properties 

119 properties = self._get_compound_properties(cid) 

120 

121 # Get compound description 

122 description = self._get_compound_description(cid) 

123 

124 return { 

125 "cid": cid, 

126 "name": name, 

127 "properties": properties, 

128 "description": description, 

129 } 

130 

131 except RateLimitError: 

132 raise 

133 except Exception as e: 

134 safe_msg = self._scrub_error(e) 

135 logger.exception( 

136 f"Error fetching PubChem compound: {name} ({type(e).__name__}): {safe_msg}" 

137 ) 

138 return None 

139 

140 def _get_compound_properties(self, cid: int) -> Dict[str, Any]: 

141 """Get properties for a compound by CID.""" 

142 try: 

143 self.rate_tracker.apply_rate_limit(self.engine_type) 

144 properties_list = [ 

145 "MolecularFormula", 

146 "MolecularWeight", 

147 "IUPACName", 

148 "CanonicalSMILES", 

149 "IsomericSMILES", 

150 "InChI", 

151 "InChIKey", 

152 "XLogP", 

153 "TPSA", 

154 "Complexity", 

155 "Charge", 

156 "HBondDonorCount", 

157 "HBondAcceptorCount", 

158 "RotatableBondCount", 

159 "HeavyAtomCount", 

160 ] 

161 

162 url = f"{self.base_url}/compound/cid/{cid}/property/{','.join(properties_list)}/JSON" 

163 response = safe_get(url, headers=self.headers, timeout=30) 

164 self._raise_if_rate_limit(response.status_code) 

165 response.raise_for_status() 

166 data = response.json() 

167 

168 props = data.get("PropertyTable", {}).get("Properties", []) 

169 return props[0] if props else {} 

170 

171 except RateLimitError: 

172 raise 

173 except Exception as e: 

174 safe_msg = self._scrub_error(e) 

175 logger.exception( 

176 f"Error fetching PubChem properties for CID {cid} ({type(e).__name__}): {safe_msg}" 

177 ) 

178 return {} 

179 

180 def _get_compound_description(self, cid: int) -> str: 

181 """Get description for a compound by CID.""" 

182 try: 

183 self.rate_tracker.apply_rate_limit(self.engine_type) 

184 url = f"{self.base_url}/compound/cid/{cid}/description/JSON" 

185 response = safe_get(url, headers=self.headers, timeout=30) 

186 

187 if response.status_code == 404: 

188 return "" 

189 self._raise_if_rate_limit(response.status_code) 

190 

191 response.raise_for_status() 

192 data = response.json() 

193 

194 descriptions = data.get("InformationList", {}).get( 

195 "Information", [] 

196 ) 

197 for desc in descriptions: 

198 if desc.get("Description"): 

199 return desc.get("Description", "") # type: ignore[no-any-return] 

200 

201 return "" 

202 

203 except RateLimitError: 

204 raise 

205 except Exception as e: 

206 safe_msg = self._scrub_error(e) 

207 logger.exception( 

208 f"Error fetching PubChem description for CID {cid} ({type(e).__name__}): {safe_msg}" 

209 ) 

210 return "" 

211 

212 def _get_compound_synonyms(self, cid: int, limit: int = 10) -> List[str]: 

213 """Get synonyms for a compound by CID.""" 

214 try: 

215 self.rate_tracker.apply_rate_limit(self.engine_type) 

216 url = f"{self.base_url}/compound/cid/{cid}/synonyms/JSON" 

217 response = safe_get(url, headers=self.headers, timeout=30) 

218 

219 if response.status_code == 404: 

220 return [] 

221 self._raise_if_rate_limit(response.status_code) 

222 

223 response.raise_for_status() 

224 data = response.json() 

225 

226 info = data.get("InformationList", {}).get("Information", []) 

227 if info: 

228 synonyms = info[0].get("Synonym", []) 

229 return synonyms[:limit] # type: ignore[no-any-return] 

230 return [] 

231 

232 except RateLimitError: 

233 raise 

234 except Exception as e: 

235 safe_msg = self._scrub_error(e) 

236 logger.exception( 

237 f"Error fetching PubChem synonyms for CID {cid} ({type(e).__name__}): {safe_msg}" 

238 ) 

239 return [] 

240 

241 def _get_previews(self, query: str) -> List[Dict[str, Any]]: 

242 """ 

243 Get preview information for PubChem compounds. 

244 

245 Args: 

246 query: The search query (compound name) 

247 

248 Returns: 

249 List of preview dictionaries 

250 """ 

251 logger.info("Getting PubChem previews") 

252 

253 # Apply rate limiting 

254 self._last_wait_time = self.rate_tracker.apply_rate_limit( 

255 self.engine_type 

256 ) 

257 

258 # Search for matching compound names 

259 compound_names = self._search_compounds(query) 

260 

261 if not compound_names: 

262 # Try direct lookup 

263 compound = self._get_compound_by_name(query) 

264 if compound: 

265 compound_names = [query] 

266 else: 

267 logger.info("No PubChem compounds found") 

268 return [] 

269 

270 logger.info(f"Found {len(compound_names)} potential compounds") 

271 

272 previews: list[dict[str, Any]] = [] 

273 seen_cids = set() 

274 for index, name in enumerate(compound_names): 

275 if len(previews) >= self.max_results: 

276 break 

277 

278 try: 

279 compound = self._get_compound_by_name(name) 

280 if not compound: 

281 continue 

282 

283 cid = compound["cid"] 

284 

285 # Deduplicate by CID (autocomplete may return 

286 # case variants like "Caffeine" and "caffeine") 

287 if cid in seen_cids: 

288 continue 

289 seen_cids.add(cid) 

290 properties = compound.get("properties", {}) 

291 description = compound.get("description", "") 

292 

293 # Build compound URL 

294 compound_url = ( 

295 f"https://pubchem.ncbi.nlm.nih.gov/compound/{cid}" 

296 ) 

297 

298 # Get key properties 

299 molecular_formula = properties.get("MolecularFormula", "") 

300 molecular_weight = properties.get("MolecularWeight", "") 

301 iupac_name = properties.get("IUPACName", "") 

302 smiles = ( 

303 properties.get("CanonicalSMILES", "") 

304 or properties.get("SMILES", "") 

305 or properties.get("IsomericSMILES", "") 

306 or properties.get("ConnectivitySMILES", "") 

307 ) 

308 

309 # Get drug-relevant properties 

310 xlogp = properties.get("XLogP") 

311 hbond_donors = properties.get("HBondDonorCount") 

312 hbond_acceptors = properties.get("HBondAcceptorCount") 

313 

314 # Build snippet 

315 snippet_parts = [] 

316 if molecular_formula: 

317 snippet_parts.append(f"Formula: {molecular_formula}") 

318 if molecular_weight: 

319 snippet_parts.append(f"MW: {molecular_weight}") 

320 if xlogp is not None: 

321 snippet_parts.append(f"XLogP: {xlogp}") 

322 if hbond_donors is not None or hbond_acceptors is not None: 

323 hbond_info = [] 

324 if hbond_donors is not None: 324 ↛ 326line 324 didn't jump to line 326 because the condition on line 324 was always true

325 hbond_info.append(f"H-Donors: {hbond_donors}") 

326 if hbond_acceptors is not None: 326 ↛ 328line 326 didn't jump to line 328 because the condition on line 326 was always true

327 hbond_info.append(f"H-Acceptors: {hbond_acceptors}") 

328 snippet_parts.append(", ".join(hbond_info)) 

329 if iupac_name: 

330 snippet_parts.append(f"IUPAC: {iupac_name}") 

331 if description: 

332 snippet_parts.append(description[:200]) 

333 snippet = ". ".join(snippet_parts) 

334 

335 preview = { 

336 "id": str(cid), 

337 "cid": cid, 

338 "title": name, 

339 "link": compound_url, 

340 "snippet": snippet, 

341 "molecular_formula": molecular_formula, 

342 "molecular_weight": molecular_weight, 

343 "iupac_name": iupac_name, 

344 "smiles": smiles, 

345 "inchi_key": properties.get("InChIKey", ""), 

346 "description": description, 

347 "source": "PubChem", 

348 "_raw": { 

349 "properties": properties, 

350 "description": description, 

351 }, 

352 } 

353 

354 previews.append(preview) 

355 

356 except RateLimitError: 

357 raise 

358 except Exception as e: 

359 # `name` is a compound name and, on the autocomplete-empty 

360 # fallback path (compound_names = [query] above), can be the 

361 # raw user query — log its position instead (#5646). 

362 safe_msg = self._scrub_error(e) 

363 logger.exception( 

364 f"Error processing PubChem compound {index + 1}/" 

365 f"{len(compound_names)} ({type(e).__name__}): {safe_msg}" 

366 ) 

367 continue 

368 

369 return previews 

370 

371 def _get_full_content( 

372 self, relevant_items: List[Dict[str, Any]] 

373 ) -> List[Dict[str, Any]]: 

374 """ 

375 Get full content for the relevant PubChem compounds. 

376 

377 Args: 

378 relevant_items: List of relevant preview dictionaries 

379 

380 Returns: 

381 List of result dictionaries with full content 

382 """ 

383 logger.info( 

384 f"Getting full content for {len(relevant_items)} PubChem compounds" 

385 ) 

386 

387 results = [] 

388 for item in relevant_items: 

389 result = item.copy() 

390 

391 cid = item.get("cid") 

392 if cid and self.include_synonyms: 392 ↛ 394line 392 didn't jump to line 394 because the condition on line 392 was never true

393 # Get synonyms 

394 synonyms = self._get_compound_synonyms(cid) 

395 result["synonyms"] = synonyms 

396 

397 raw = item.get("_raw", {}) 

398 if raw: 398 ↛ 450line 398 didn't jump to line 450 because the condition on line 398 was always true

399 properties = raw.get("properties", {}) 

400 description = raw.get("description", "") 

401 

402 # Build content summary 

403 content_parts = [] 

404 content_parts.append( 

405 f"Compound: {result.get('title', 'Unknown')}" 

406 ) 

407 if cid is not None: 407 ↛ 410line 407 didn't jump to line 410 because the condition on line 407 was always true

408 content_parts.append(f"CID: {cid}") 

409 

410 if result.get("molecular_formula"): 410 ↛ 414line 410 didn't jump to line 414 because the condition on line 410 was always true

411 content_parts.append( 

412 f"Molecular Formula: {result['molecular_formula']}" 

413 ) 

414 if result.get("molecular_weight"): 414 ↛ 418line 414 didn't jump to line 418 because the condition on line 414 was always true

415 content_parts.append( 

416 f"Molecular Weight: {result['molecular_weight']} g/mol" 

417 ) 

418 if result.get("iupac_name"): 418 ↛ 420line 418 didn't jump to line 420 because the condition on line 418 was always true

419 content_parts.append(f"IUPAC Name: {result['iupac_name']}") 

420 if result.get("smiles"): 420 ↛ 422line 420 didn't jump to line 422 because the condition on line 420 was always true

421 content_parts.append(f"SMILES: {result['smiles']}") 

422 if result.get("inchi_key"): 422 ↛ 426line 422 didn't jump to line 426 because the condition on line 422 was always true

423 content_parts.append(f"InChIKey: {result['inchi_key']}") 

424 

425 # Additional properties 

426 if properties.get("XLogP") is not None: 426 ↛ 428line 426 didn't jump to line 428 because the condition on line 426 was always true

427 content_parts.append(f"XLogP: {properties['XLogP']}") 

428 if properties.get("TPSA") is not None: 428 ↛ 430line 428 didn't jump to line 430 because the condition on line 428 was always true

429 content_parts.append(f"TPSA: {properties['TPSA']} Ų") 

430 if properties.get("HBondDonorCount") is not None: 430 ↛ 434line 430 didn't jump to line 434 because the condition on line 430 was always true

431 content_parts.append( 

432 f"H-Bond Donors: {properties['HBondDonorCount']}" 

433 ) 

434 if properties.get("HBondAcceptorCount") is not None: 434 ↛ 439line 434 didn't jump to line 439 because the condition on line 434 was always true

435 content_parts.append( 

436 f"H-Bond Acceptors: {properties['HBondAcceptorCount']}" 

437 ) 

438 

439 if result.get("synonyms"): 439 ↛ 440line 439 didn't jump to line 440 because the condition on line 439 was never true

440 content_parts.append( 

441 f"\nSynonyms: {', '.join(result['synonyms'][:5])}" 

442 ) 

443 

444 if description: 444 ↛ 447line 444 didn't jump to line 447 because the condition on line 444 was always true

445 content_parts.append(f"\nDescription: {description}") 

446 

447 result["content"] = "\n".join(content_parts) 

448 

449 # Clean up internal fields 

450 if "_raw" in result: 450 ↛ 453line 450 didn't jump to line 453 because the condition on line 450 was always true

451 del result["_raw"] 

452 

453 results.append(result) 

454 

455 return results 

456 

457 def get_compound(self, cid: int) -> Optional[Dict[str, Any]]: 

458 """ 

459 Get a specific compound by CID. 

460 

461 Args: 

462 cid: The PubChem compound ID 

463 

464 Returns: 

465 Compound dictionary or None 

466 """ 

467 try: 

468 properties = self._get_compound_properties(cid) 

469 description = self._get_compound_description(cid) 

470 synonyms = self._get_compound_synonyms(cid) 

471 

472 return { 

473 "cid": cid, 

474 "properties": properties, 

475 "description": description, 

476 "synonyms": synonyms, 

477 } 

478 except RateLimitError: 

479 raise 

480 except Exception as e: 

481 safe_msg = self._scrub_error(e) 

482 logger.exception( 

483 f"Error fetching PubChem compound {cid} ({type(e).__name__}): {safe_msg}" 

484 ) 

485 return None 

486 

487 def search_by_formula(self, formula: str) -> List[Dict[str, Any]]: 

488 """ 

489 Search compounds by molecular formula. 

490 

491 Args: 

492 formula: Molecular formula (e.g., "C6H12O6") 

493 

494 Returns: 

495 List of matching compounds 

496 """ 

497 try: 

498 url = f"{self.base_url}/compound/fastformula/{quote(formula, safe='')}/cids/JSON" 

499 response = safe_get(url, headers=self.headers, timeout=30) 

500 

501 if response.status_code == 404: 

502 return [] 

503 self._raise_if_rate_limit(response.status_code) 

504 

505 response.raise_for_status() 

506 data = response.json() 

507 cids = data.get("IdentifierList", {}).get("CID", []) 

508 

509 results = [] 

510 for cid in cids[: self.max_results]: 

511 compound = self.get_compound(cid) 

512 if compound: 

513 results.append(compound) 

514 

515 return results 

516 

517 except RateLimitError: 

518 raise 

519 except Exception as e: 

520 safe_msg = self._scrub_error(e) 

521 logger.exception( 

522 f"Error searching by formula: {formula} ({type(e).__name__}): {safe_msg}" 

523 ) 

524 return []