import os import json import logging import asyncio import threading from typing import Optional, Callable, Coroutine, Any from src.config import settings from src.models.paper import Paper from src.models.claim import ExtractedClaim, ClaimExtractionResponse from src.llm.base import LLMProvider logger = logging.getLogger(__name__) _section_semaphores = {} _section_semaphore_lock = threading.Lock() def get_section_semaphore() -> asyncio.Semaphore: try: loop = asyncio.get_running_loop() except RuntimeError: return asyncio.Semaphore(settings.section_concurrency) with _section_semaphore_lock: # Clean up any closed loops to prevent memory leaks and loop ID collision issues for old_loop in list(_section_semaphores.keys()): if old_loop.is_closed(): del _section_semaphores[old_loop] if loop not in _section_semaphores: _section_semaphores[loop] = asyncio.Semaphore(settings.section_concurrency) return _section_semaphores[loop] # Cache prompt templates to avoid repeated disk reads _PROMPT_TEMPLATE = None _FEW_SHOTS_DATA = None def load_prompt_resources() -> tuple[str, str]: """Load prompt template and few-shot examples from files.""" global _PROMPT_TEMPLATE, _FEW_SHOTS_DATA if _PROMPT_TEMPLATE is None: base_dir = os.path.dirname(os.path.abspath(__file__)) prompt_path = os.path.join(base_dir, "prompts", "extraction_prompt.txt") try: with open(prompt_path, "r", encoding="utf-8") as f: _PROMPT_TEMPLATE = f.read() except FileNotFoundError: logger.error(f"Prompt file not found: {prompt_path}") _PROMPT_TEMPLATE = "Extract claims from the following text:" few_shots_path = os.path.join(base_dir, "prompts", "extraction_few_shot.json") try: with open(few_shots_path, "r", encoding="utf-8") as f: raw_shots = json.load(f) shots_str = "" for idx, shot in enumerate(raw_shots): shots_str += f"### Example {idx+1}\n" shots_str += f"Abstract: {shot['abstract']}\n" formatted_claims = json.dumps({"claims": shot["claims"]}, indent=2) shots_str += f"Extracted Output:\n```json\n{formatted_claims}\n```\n\n" _FEW_SHOTS_DATA = shots_str except FileNotFoundError: logger.warning(f"Few-shot file not found: {few_shots_path}. Using empty few-shot examples.") _FEW_SHOTS_DATA = "" return _PROMPT_TEMPLATE, _FEW_SHOTS_DATA def build_extraction_prompt(text: str, is_full_text: bool = False, section_name: str | None = None) -> str: """Combine the base prompt template, formatted few-shot examples, and target text.""" template, few_shots = load_prompt_resources() if section_name: text_type = f"Section '{section_name}'" else: text_type = "Full Text (including abstract)" if is_full_text else "Abstract" # Adapt template slightly if full text or section is used if is_full_text or section_name: template = template.replace("research paper abstract", "research paper text") template = template.replace("substring of the abstract", "substring of the text") template = template.replace("claims per abstract", "claims per paper") full_prompt = ( f"{template}\n" f"Here are reference examples of the expected input/output mapping:\n\n" f"{few_shots}\n" f"Please extract claims for this target {text_type.lower()}:\n" f"{text_type}: {text}\n\n" f"Extracted Output (JSON format matching ClaimExtractionResponse schema):" ) return full_prompt async def extract_claims_from_paper( paper: Paper, llm: LLMProvider, ) -> list[ExtractedClaim]: """Extract claims from a single paper's text (full text if available, fallback to abstract). If full text is available: - Splits it into sections (Introduction, Methods, Results, Discussion). - Concurrently extracts claims from each non-empty section. - Caps claims per section to prevent any single section from dominating. - Merges results and caps final claims list at claims_per_abstract_cap. If no sections are found or full text is not available: - Falls back to extracting from abstract text / full text as a single string. """ import re # Check if we have sections in full text to do section-by-section extraction sections_to_extract = [] if paper.full_text: parts = re.split(r'===\s*([A-Z\s]+)\s*===', paper.full_text) primary_names = {s.lower() for s in settings.primary_section_names} for i in range((len(parts) - 1) // 2): sec_header = parts[i * 2 + 1].strip() sec_content = parts[i * 2 + 2].strip() if len(sec_content) <= 100: continue if settings.primary_sections_only and sec_header.lower() not in primary_names: logger.debug( f"Skipping section '{sec_header}' for paper {paper.pmid} " f"(not in primary_section_names, primary_sections_only=True)" ) continue sections_to_extract.append((sec_header, sec_content)) if sections_to_extract: # Run section-based extraction concurrently async def extract_sec(sec_title, sec_body): prompt = build_extraction_prompt(sec_body, section_name=sec_title) sem = get_section_semaphore() async with sem: for attempt in range(2): try: response = await llm.generate_structured( prompt=prompt, response_schema=ClaimExtractionResponse, temperature=0.1 ) sec_claims = response.claims sec_cap = max(3, settings.claims_per_abstract_cap // 2) if len(sec_claims) > sec_cap: sec_claims = sec_claims[:sec_cap] return sec_claims except Exception as e: if attempt == 0: logger.warning(f"Attempt 1 failed to extract from section {sec_title} of paper {paper.pmid}, retrying: {e}") await asyncio.sleep(1) else: logger.error(f"Failed to extract from section {sec_title} of paper {paper.pmid} after 2 attempts: {e}") return [] tasks = [extract_sec(title, content) for title, content in sections_to_extract] results = await asyncio.gather(*tasks) merged_claims = [] for sec_claims in results: merged_claims.extend(sec_claims) if len(merged_claims) > settings.claims_per_abstract_cap: logger.info(f"Merged claims count for full-text paper {paper.pmid} exceeded cap ({len(merged_claims)}). Capping at {settings.claims_per_abstract_cap}.") merged_claims = merged_claims[:settings.claims_per_abstract_cap] return merged_claims # Fallback to single-call abstract/full-text extraction source_text = paper.full_text or paper.abstract_text is_full_text = bool(paper.full_text) if not source_text or not source_text.strip(): logger.warning(f"Paper {paper.pmid} has empty abstract and full text, skipping claim extraction.") return [] prompt = build_extraction_prompt(source_text, is_full_text=is_full_text) for attempt in range(2): try: response = await llm.generate_structured( prompt=prompt, response_schema=ClaimExtractionResponse, temperature=0.1 ) claims = response.claims if len(claims) > settings.claims_per_abstract_cap: logger.info(f"Claims count for {paper.pmid} exceeded cap ({len(claims)}). Capping at {settings.claims_per_abstract_cap}.") claims = claims[:settings.claims_per_abstract_cap] return claims except Exception as e: if attempt == 0: logger.warning(f"Attempt 1 failed to extract claims from paper {paper.pmid}, retrying: {e}") await asyncio.sleep(1) else: logger.error(f"Failed to extract claims from paper {paper.pmid} after 2 attempts: {e}") return [] async def extract_claims_batch( papers: list[Paper], llm: LLMProvider, on_paper_complete: Optional[Callable[[Paper, list[ExtractedClaim]], Coroutine[Any, Any, None]]] = None ) -> dict[str, list[ExtractedClaim]]: """Extract claims from all papers concurrently. Uses asyncio.Semaphore(settings.llm_concurrency) for rate limiting. Returns: {pmid: [claims]} mapping Logs: papers that failed extraction, papers with 0 claims """ sem = asyncio.Semaphore(settings.llm_concurrency) async def process_paper(paper: Paper): async with sem: claims = await extract_claims_from_paper(paper, llm) if on_paper_complete: await on_paper_complete(paper, claims) return paper.pmid, claims tasks = [process_paper(p) for p in papers] results = await asyncio.gather(*tasks) extracted = {} total_claims = 0 for pmid, claims in results: extracted[pmid] = claims total_claims += len(claims) if len(claims) == 0: logger.info(f"Paper {pmid} returned 0 claims during extraction.") logger.info(f"Batch extraction completed. Extracted {total_claims} claims from {len(papers)} papers.") return extracted