"""LangGraph orchestration: search -> filter -> read -> score -> synthesize (-> deepen).""" from __future__ import annotations import json import re from concurrent.futures import ThreadPoolExecutor from typing import Callable, TypedDict from . import utils from .db import PaperVectorDB from .tools import ( DownloadTool, ExtractionTool, Paper, SearchTool, SemanticScholarTool, ) class AgentState(TypedDict): topic: str papers: list[Paper] review: str gaps: list[str] depth: int iteration: int NUM_QUERIES = 3 RESULTS_PER_QUERY = 8 DEFAULT_MAX_PAPERS = 15 POOL_BUFFER = 3 READ_WORKERS = 4 # parallel PDF download + extraction def _generate_queries(topic: str, meter: dict | None = None, model: str | None = None) -> list[str]: prompt = ( f"Generate {NUM_QUERIES} diverse, specific search queries that together " f'cover the research topic: "{topic}".\n' "Vary the angle (methods, applications, theory, recent advances). " "Return ONLY a JSON array of short query strings." ) try: raw = utils.complete(prompt, meter=meter, model=model) match = re.search(r"\[.*\]", raw, flags=re.DOTALL) queries = [str(q).strip() for q in (json.loads(match.group(0)) if match else []) if str(q).strip()] except Exception: queries = [] if topic not in queries: queries.insert(0, topic) return queries[:NUM_QUERIES] or [topic] def _queries_from_gaps(topic, gaps, meter=None, model=None) -> list[str]: gap_text = "\n".join(f"- {g}" for g in gaps[:5]) prompt = ( f'For the topic "{topic}", here are open problems/gaps identified so far:\n' f"{gap_text}\n\nGenerate {NUM_QUERIES} new, specific search queries to find " "papers addressing these gaps. Return ONLY a JSON array of short query strings." ) try: raw = utils.complete(prompt, meter=meter, model=model) match = re.search(r"\[.*\]", raw, flags=re.DOTALL) return [str(q).strip() for q in (json.loads(match.group(0)) if match else []) if str(q).strip()][:NUM_QUERIES] except Exception: return [] def _filter_relevant(topic, papers, keep, meter=None, model=None) -> list[Paper]: """Drop clearly off-topic candidates with one cheap LLM relevance pass.""" if len(papers) <= keep: return papers listing = "\n".join( f"[{i}] {p.title} — {p.abstract[:160]}" for i, p in enumerate(papers) ) prompt = ( f'Topic: "{topic}"\n\nCandidate papers:\n{listing}\n\n' f"Return ONLY a JSON array of the indices (numbers) of the {keep} papers most " "relevant to the topic, best first. Exclude clearly off-topic papers." ) try: raw = utils.complete(prompt, meter=meter, model=model) match = re.search(r"\[.*\]", raw, flags=re.DOTALL) idxs = json.loads(match.group(0)) if match else [] chosen = [papers[i] for i in idxs if isinstance(i, int) and 0 <= i < len(papers)] return chosen or papers[:keep] except Exception: return papers[:keep] def build_graph( max_papers: int = DEFAULT_MAX_PAPERS, progress: Callable[[str], None] | None = None, meter: dict | None = None, model: str | None = None, year_min: int = 0, style: str = "concise", extract_model: str | None = None, ): """Build and compile the research agent graph. ``style`` is "concise" (default — shorter review, far cheaper since output tokens dominate cost) or "comprehensive" (longer, more detailed). ``extract_model`` overrides the model for the ~20 mechanical extraction calls (e.g. a cheaper flash-lite), keeping the run model for synthesis. """ from langgraph.graph import END, StateGraph say = progress or (lambda _msg: None) searcher = SearchTool(max_results=RESULTS_PER_QUERY) scholar = SemanticScholarTool(max_results=RESULTS_PER_QUERY, year_min=year_min) downloader = DownloadTool() extractor = ExtractionTool(model=extract_model or model, meter=meter) pool_cap = max_papers + POOL_BUFFER def _run_searches(queries: list[str], existing: list[Paper]) -> list[Paper]: merged: dict[str, Paper] = {p.id: p for p in existing} seen_titles = {p.title.lower() for p in existing} def add(papers: list[Paper]) -> None: for p in papers: if p.id in merged or p.title.lower() in seen_titles: continue merged[p.id] = p seen_titles.add(p.title.lower()) for q in queries: say(f" query: {q}") try: add(searcher.search(q)) except Exception as err: say(f" (arXiv query failed: {err})") try: add(scholar.search(q)) except Exception as err: say(f" (Semantic Scholar query failed: {err})") return list(merged.values()) def search_node(state: AgentState) -> dict: topic = state["topic"] queries = _generate_queries(topic, meter=meter, model=model) say(f"Searching arXiv + Semantic Scholar with {len(queries)} queries...") candidates = _run_searches(queries, state.get("papers", [])) n_s2 = sum(1 for p in candidates if p.source == "semantic_scholar") say(f"Found {len(candidates)} candidates ({len(candidates) - n_s2} arXiv, {n_s2} S2).") # Relevance filter: drop off-topic papers before the expensive read step. already_read = [p for p in candidates if p.full_text] fresh = [p for p in candidates if not p.full_text] if fresh: say("Filtering candidates by relevance...") fresh = _filter_relevant(topic, fresh, pool_cap, meter=meter, model=model) papers = (already_read + fresh)[: pool_cap + state.get("iteration", 0) * POOL_BUFFER] say(f"Keeping {len(papers)} relevant papers.") return {"papers": papers} def read_node(state: AgentState) -> dict: papers = state["papers"] if not papers: raise RuntimeError( "No papers found from arXiv or Semantic Scholar — both sources may " "be rate-limiting right now. Please try again in a minute." ) # Only read the top max_papers (already ranked by score_node); the rest # were filtered/ranked on abstracts alone, so we never pay to read them. todo = [p for p in papers[:max_papers] if not p.full_text] say(f"Reading top {len(todo)} papers (parallel)...") def _read_one(paper: Paper) -> None: paper.full_text = downloader.get_text(paper) try: extractor.extract(paper) except Exception as err: say(f" (extraction failed: {paper.title[:40]}: {err})") tag = "abstract-only" if paper.read_from == "abstract" else "pdf" say(f" read: {paper.title[:55]} [{tag}]") with ThreadPoolExecutor(max_workers=READ_WORKERS) as pool: list(pool.map(_read_one, todo)) return {"papers": papers} def score_node(state: AgentState) -> dict: papers = state["papers"] say("Scoring papers by relevance...") db = None try: db = PaperVectorDB() db.index(papers) ranked_ids = db.search(state["topic"], n=len(papers)) order = {pid: rank for rank, pid in enumerate(ranked_ids)} for p in papers: p.score = 1.0 / (1 + order.get(p.id, len(papers))) papers.sort(key=lambda p: order.get(p.id, len(papers))) except Exception as err: say(f" (scoring skipped: {err})") finally: if db is not None: db.close() return {"papers": papers} def synthesize_node(state: AgentState) -> dict: top = state["papers"][:max_papers] say(f"Synthesizing {style} review from top {len(top)} papers...") papers_block = utils.format_papers_for_synthesis(top) if style == "comprehensive": length_note = "Write a thorough, detailed review." section_note = ( "2. ## Key Themes (group by methodology)\n" "3. ## Key Findings\n" "4. ## Open Problems & Gaps\n" ) else: # concise (default): brevity comes from the prompt, not a token cap, # so the review never truncates — it just stays short (and cheap). length_note = ( "Write a CONCISE review — 600 words MAXIMUM for the whole thing. " "Keep EVERY section to 2-4 tight sentences. Do NOT write a paragraph " "or bullet per paper or per theme; summarize across them. Prefer " "brevity over completeness. Still include all five sections below." ) section_note = ( "2. ## Key Themes (2-4 sentences across methodologies)\n" "3. ## Key Findings (2-4 sentences)\n" "4. ## Open Problems & Gaps (2-4 sentences)\n" ) prompt = ( f"{length_note}\n\nLiterature review on: {state['topic']}\n\n" f"Papers:\n{papers_block}\n\n" "Use these Markdown sections:\n" "1. ## Introduction\n" f"{section_note}" "5. ## Summary Table (Markdown table: Paper | Year | Method | Key result)\n\n" "Be specific and grounded ONLY in the papers above. Every claim must cite " "its source as [Author, Year] using the first author's surname. Do not " "invent papers, authors, or findings not present above." ) # High ceiling for BOTH styles so nothing ever truncates; concise stays # cheap because the prompt keeps the actual output short. review = utils.complete(prompt, max_tokens=16384, meter=meter, model=model) review = review.rstrip() + "\n" + utils.references_markdown(top) gaps: list[str] = [] section = re.search( r"##\s*Open Problems.*?\n(.*?)(?:\n##\s|\Z)", review, flags=re.DOTALL ) if section: gaps = [ line.strip("-* ").strip() for line in section.group(1).splitlines() if line.strip().startswith(("-", "*")) ] return {"review": review, "gaps": gaps} def deepen_node(state: AgentState) -> dict: it = state.get("iteration", 0) + 1 say(f"Deepening (pass {it + 1}/{state['depth']}): searching gaps...") queries = _queries_from_gaps(state["topic"], state.get("gaps", []), meter=meter, model=model) if not queries: return {"iteration": it} candidates = _run_searches(queries, state["papers"]) fresh = [p for p in candidates if not p.full_text] kept_fresh = _filter_relevant(state["topic"], fresh, pool_cap, meter=meter, model=model) if fresh else [] papers = [p for p in candidates if p.full_text] + kept_fresh say(f"Now {len(papers)} candidate papers after deepening.") return {"papers": papers, "iteration": it} def should_deepen(state: AgentState) -> str: if state.get("iteration", 0) < state.get("depth", 1) - 1 and state.get("gaps"): return "deepen" return END graph = StateGraph(AgentState) graph.add_node("search", search_node) graph.add_node("read", read_node) graph.add_node("score", score_node) graph.add_node("synthesize", synthesize_node) graph.add_node("deepen", deepen_node) graph.set_entry_point("search") graph.add_edge("search", "score") # rank on abstracts first... graph.add_edge("score", "read") # ...then read only the top max_papers graph.add_edge("read", "synthesize") graph.add_conditional_edges("synthesize", should_deepen, {"deepen": "deepen", END: END}) graph.add_edge("deepen", "score") return graph.compile()