src/agent.py

"""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()