Aminer Deep Search

from __future__ import annotations

import math
import os
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Iterable, Sequence

import requests


AMINER_PAPER_DETAIL_URL = "https://datacenter.aminer.cn/gateway/api/v3/paper/detail/batch"


def get_aminer_key() -> str:
    aminer_key = os.getenv("AMINER_API_KEY") or os.getenv("AMINER_KEY")
    if not aminer_key:
        raise ValueError("AMINER_API_KEY is required for AMiner API calls.")
    return aminer_key


def json_auth_headers() -> dict[str, str]:
    return {
        "Content-Type": "application/json;charset=utf-8",
        "Authorization": f"Bearer {get_aminer_key()}",
    }


def dedupe_preserve_order(items: Iterable[str]) -> list[str]:
    deduped: list[str] = []
    seen: set[str] = set()
    for item in items:
        cleaned = str(item).strip()
        if not cleaned or cleaned in seen:
            continue
        seen.add(cleaned)
        deduped.append(cleaned)
    return deduped


def chunks(items: Sequence[Any], chunk_size: int) -> Iterable[Sequence[Any]]:
    for index in range(0, len(items), chunk_size):
        yield items[index : index + chunk_size]


def extract_paper_id(detail: Any) -> str:
    if isinstance(detail, dict):
        return str(detail.get("id") or detail.get("_id") or "")
    if detail is None:
        return ""
    return str(detail)


def safe_int(value: Any, default: int = 0) -> int:
    try:
        if value is None:
            return default
        return int(value)
    except (TypeError, ValueError):
        return default


def normalize_authors(authors: Any) -> list[str]:
    if not isinstance(authors, list):
        return []
    normalized: list[str] = []
    for author in authors:
        if isinstance(author, dict):
            name = author.get("name") or author.get("name_zh")
        else:
            name = str(author)
        if name:
            normalized.append(str(name))
    return normalized


def normalize_paper_detail(detail: dict[str, Any], *, query: str = "") -> dict[str, Any]:
    venue = detail.get("venue")
    if isinstance(venue, dict):
        venue_text = venue.get("raw") or venue.get("name") or ""
    else:
        venue_text = venue or detail.get("venue_name") or ""

    orgs = detail.get("orgs") or detail.get("organizations") or detail.get("affiliations") or []
    if isinstance(orgs, str):
        organizations = [orgs]
    elif isinstance(orgs, list):
        organizations = [str(org) for org in orgs if org]
    else:
        organizations = []

    normalized = dict(detail)
    normalized["id"] = extract_paper_id(detail)
    normalized["title"] = detail.get("title") or detail.get("title_zh") or ""
    normalized["abstract"] = detail.get("abstract") or detail.get("abstract_zh") or ""
    normalized["authors"] = normalize_authors(detail.get("authors"))
    normalized["organization"] = organizations
    normalized["venue"] = str(venue_text)
    normalized["year"] = detail.get("year")
    normalized["n_citation"] = safe_int(detail.get("n_citation") or detail.get("num_citation"), 0)
    normalized["keywords"] = detail.get("keywords") or []
    normalized["score"] = round(rule_based_score(normalized, query=query), 4)
    return normalized


def tokenize(text: str) -> list[str]:
    return re.findall(r"[A-Za-z0-9][A-Za-z0-9_-]{1,}", text.lower())


def rule_based_score(paper: dict[str, Any], *, query: str = "") -> float:
    query_tokens = set(tokenize(query))
    title = str(paper.get("title") or "")
    abstract = str(paper.get("abstract") or "")
    keywords = " ".join(str(item) for item in paper.get("keywords") or [])
    haystack = f"{title} {abstract} {keywords}"
    haystack_tokens = set(tokenize(haystack))

    lexical = 0.0
    if query_tokens:
        lexical = len(query_tokens & haystack_tokens) / max(1, len(query_tokens))

    phrase_bonus = 0.2 if query and query.lower() in haystack.lower() else 0.0
    citation_score = min(0.3, math.log1p(safe_int(paper.get("n_citation"), 0)) / 30.0)
    year = safe_int(paper.get("year"), 0)
    recency_score = 0.1 if year >= 2020 else 0.05 if year >= 2015 else 0.0
    return min(1.0, lexical * 0.55 + phrase_bonus + citation_score + recency_score)


def request_paper_detail_batch(paper_ids: Sequence[str]) -> list[dict[str, Any]]:
    ids = dedupe_preserve_order(paper_ids)
    if not ids:
        return []
    try:
        response = requests.post(
            AMINER_PAPER_DETAIL_URL,
            json={"ids": ids},
            headers=json_auth_headers(),
            timeout=(10, 30),
        )
        if response.status_code != 200:
            print(f"AMiner detail request failed: status={response.status_code}, detail={response.text[:300]}")
            return []
        data = response.json().get("data", [])
    except (requests.RequestException, ValueError) as exc:
        print(f"AMiner detail request failed: {exc}")
        return []

    if isinstance(data, dict):
        data = data.get("data") or data.get("items") or []
    if not isinstance(data, list):
        return []
    return [item for item in data if isinstance(item, dict)]


def aminer_get_paper_info_batch(
    paper_ids: Sequence[str],
    detail_batch_size: int = 50,
    max_workers: int = 8,
) -> list[dict[str, Any]]:
    ids = dedupe_preserve_order(paper_ids)
    if not ids:
        return []
    batches = list(chunks(ids, max(1, int(detail_batch_size))))
    if len(batches) == 1:
        return request_paper_detail_batch(batches[0])

    results_by_index: dict[int, list[dict[str, Any]]] = {}
    with ThreadPoolExecutor(max_workers=max(1, min(max_workers, len(batches)))) as executor:
        future_to_index = {
            executor.submit(request_paper_detail_batch, batch): index
            for index, batch in enumerate(batches)
        }
        for future in as_completed(future_to_index):
            index = future_to_index[future]
            try:
                results_by_index[index] = future.result()
            except Exception as exc:
                print(f"Failed to fetch AMiner detail batch: {exc}")
                results_by_index[index] = []

    details: list[dict[str, Any]] = []
    for index in range(len(batches)):
        details.extend(results_by_index.get(index, []))
    return details

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Sequence

from openai import OpenAI


@dataclass
class Response:
    """API response wrapper."""

    content: str
    reasoning_content: str = ""


class APIClient:
    """OpenAI-compatible chat client with model fallback."""

    MODEL_NAME_LIST: list[str] = []

    def __init__(
        self,
        api_key: str | None,
        base_url: str | None = None,
        timeout: float = 20,
    ) -> None:
        if not api_key:
            raise ValueError(
                "LLM API key is required. Set env LLM_API_KEY (legacy: llm.api_key) or pass --api-key."
            )

        self.base_url = base_url
        self.api_key = api_key
        self.timeout = timeout
        client_kwargs: dict[str, Any] = {"api_key": self.api_key, "timeout": timeout}
        if self.base_url:
            client_kwargs["base_url"] = self.base_url
        self.client = OpenAI(**client_kwargs)

    def call(
        self,
        message: str,
        system_prompt: str | None = None,
        model_list: Sequence[str] | None = None,
        validator: Callable[[str], bool] | None = None,
    ) -> tuple[Response, bool]:
        messages: list[dict[str, str]] = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": message})
        return self.call_messages(messages, model_list=model_list, validator=validator)

    def call_messages(
        self,
        messages: Sequence[dict[str, Any]],
        model_list: Sequence[str] | None = None,
        validator: Callable[[str], bool] | None = None,
    ) -> tuple[Response, bool]:
        models = list(model_list or self.MODEL_NAME_LIST)
        if not models:
            raise ValueError("LLM model is required. Set env LLM_MODEL (legacy: llm.model) or pass --models.")
        for model in models:
            try:
                print(f"Trying model: {model}")
                response = self.client.chat.completions.create(
                    model=model,
                    messages=list(messages),
                    timeout=self.timeout,
                )
                message = response.choices[0].message
                content = message.content or ""
                reasoning = getattr(message, "reasoning_content", "") or ""
                if validator is not None and not validator(content):
                    print(f"Model {model} returned unparsable content; trying next model.")
                    continue
                return Response(content=content, reasoning_content=reasoning), True
            except Exception as exc:
                print(f"Model {model} failed: {exc}; trying next model.")
                continue
        print("All model calls failed.")
        return Response(content="", reasoning_content=""), False

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Sequence

import requests

from _utils import (
    aminer_get_paper_info_batch,
    dedupe_preserve_order,
    extract_paper_id,
    get_aminer_key,
    normalize_paper_detail,
    safe_int,
)


AMINER_CITATION_URL = "https://datacenter.aminer.cn/gateway/api/v3/paper/pub_relation"


def _auth_headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {get_aminer_key()}"}


def _fetch_pub_relation(params: dict[str, Any]) -> list[dict[str, Any]]:
    try:
        response = requests.get(
            AMINER_CITATION_URL,
            params=params,
            headers=_auth_headers(),
            timeout=(10, 30),
        )
        if response.status_code != 200:
            print(f"AMiner reference request failed: status={response.status_code}, detail={response.text[:300]}")
            return []
        data = response.json().get("data", [])
    except (requests.RequestException, ValueError) as exc:
        print(f"AMiner reference request failed: {exc}")
        return []
    return data if isinstance(data, list) else []


def fetch_references(paper_id: str, *, size: int = 20) -> list[str]:
    relations = _fetch_pub_relation({"ref": paper_id, "offset": 0, "size": max(1, int(size))})
    return dedupe_preserve_order(
        str(item.get("cited") or "").strip()
        for item in relations
        if isinstance(item, dict) and item.get("cited")
    )


def fetch_related_papers(paper_id: str, *, size: int = 20) -> list[str]:
    reference_ids = fetch_references(paper_id, size=size)
    cited_by_relations = _fetch_pub_relation({"cited": paper_id, "offset": 0, "size": max(1, int(size))})
    citing_ids = [
        str(item.get("ref") or "").strip()
        for item in cited_by_relations
        if isinstance(item, dict) and item.get("ref")
    ]
    return dedupe_preserve_order([*reference_ids, *citing_ids])


def get_reference_papers(
    aminer_ids: Sequence[str],
    *,
    topic: str = "",
    size_per_paper: int = 20,
    include_citing: bool = False,
    max_workers: int = 8,
) -> list[dict[str, Any]]:
    seed_ids = dedupe_preserve_order(aminer_ids)
    if not seed_ids:
        return []

    id_to_sources: dict[str, set[str]] = {}
    ordered_ids: list[str] = []
    seen: set[str] = set(seed_ids)

    fetcher = fetch_related_papers if include_citing else fetch_references
    with ThreadPoolExecutor(max_workers=max(1, min(max_workers, len(seed_ids)))) as executor:
        future_to_seed = {
            executor.submit(fetcher, seed_id, size=size_per_paper): seed_id
            for seed_id in seed_ids
        }
        for future in as_completed(future_to_seed):
            seed_id = future_to_seed[future]
            try:
                related_ids = future.result()
            except Exception as exc:
                print(f"Failed to fetch references for `{seed_id}`: {exc}")
                continue
            for paper_id in related_ids:
                if not paper_id or paper_id in seen:
                    continue
                seen.add(paper_id)
                ordered_ids.append(paper_id)
                id_to_sources.setdefault(paper_id, set()).add(seed_id)

    details = aminer_get_paper_info_batch(ordered_ids)
    detail_by_id = {
        extract_paper_id(detail): detail
        for detail in details
        if extract_paper_id(detail)
    }

    papers: list[dict[str, Any]] = []
    for paper_id in ordered_ids:
        detail = detail_by_id.get(paper_id)
        if not detail:
            continue
        normalized = normalize_paper_detail(detail, query=topic)
        normalized["source_paper_ids"] = sorted(id_to_sources.get(paper_id, []))
        if normalized["id"] and normalized["title"]:
            papers.append(normalized)

    papers.sort(key=lambda item: (float(item.get("score", 0.0)), safe_int(item.get("n_citation"), 0)), reverse=True)
    return papers


def citation_adding(
    total_paper_details: Sequence[Any] | None,
    uncited_paper_details: Sequence[Any] | None,
    topic: str,
    **kwargs: Any,
) -> list[dict[str, Any]]:
    existing_ids = {
        extract_paper_id(item)
        for item in (total_paper_details or [])
        if extract_paper_id(item)
    }
    seed_ids = [
        extract_paper_id(item)
        for item in (uncited_paper_details or [])
        if extract_paper_id(item)
    ]
    papers = get_reference_papers(seed_ids, topic=topic, **kwargs)
    return [paper for paper in papers if paper["id"] not in existing_ids]


citations_adding = citation_adding


__all__ = [
    "citation_adding",
    "citations_adding",
    "fetch_references",
    "fetch_related_papers",
    "get_reference_papers",
]

name: aminer-deep-search
channels:
  - conda-forge
dependencies:
  - python=3.11
  - pip
  - pip:
      - -r requirements.txt

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Iterable

from _utils import aminer_get_paper_info_batch, normalize_paper_detail


def paper_id_of(paper: Any) -> str:
    if isinstance(paper, dict):
        return str(paper.get("id") or paper.get("_id") or "")
    return str(paper or "")


class PaperSet:
    def __init__(self) -> None:
        self._papers: dict[str, dict[str, Any]] = {}

    def __len__(self) -> int:
        return len(self._papers)

    def add_papers(self, papers: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
        for paper in papers:
            paper_id = paper_id_of(paper).strip()
            title = str(paper.get("title") or paper.get("title_zh") or "").strip()
            if not paper_id or not title:
                continue
            if paper_id in self._papers:
                self._papers[paper_id].update({k: v for k, v in paper.items() if v not in (None, "", [])})
            else:
                self._papers[paper_id] = dict(paper)
        return self.all_papers()

    def add_ids(self, ids: Iterable[str], paper_cache: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
        missing_ids: list[str] = []
        papers: list[dict[str, Any]] = []
        for paper_id in ids:
            cleaned = str(paper_id).strip()
            if not cleaned:
                continue
            if cleaned in paper_cache:
                papers.append(paper_cache[cleaned])
            else:
                missing_ids.append(cleaned)

        if missing_ids:
            for detail in aminer_get_paper_info_batch(missing_ids):
                normalized = normalize_paper_detail(detail)
                if normalized["id"]:
                    paper_cache[normalized["id"]] = normalized
                    papers.append(normalized)

        return self.add_papers(papers)

    def all_papers(self) -> list[dict[str, Any]]:
        return sorted(
            self._papers.values(),
            key=lambda item: (float(item.get("score", 0.0) or 0.0), int(item.get("n_citation") or 0)),
            reverse=True,
        )

    def output(self, *, include_abstracts: bool = False) -> list[dict[str, Any]]:
        papers: list[dict[str, Any]] = []
        for paper in self.all_papers():
            item = {"id": paper["id"], "title": paper["title"]}
            if include_abstracts and paper.get("abstract"):
                item["abstract"] = paper["abstract"]
            papers.append(item)
        return papers

    def save(self, path: Path, *, include_abstracts: bool = False) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as file:
            json.dump(self.output(include_abstracts=include_abstracts), file, ensure_ascii=False, indent=2)

REACT_SYSTEM_PROMPT = """You are an academic assistant. Your task is to comprehensively collect papers related to a user-specified research topic through multiple rounds of automatic tool calls, and finally output a paper list to support writing a survey.

Objectives:

1. Use the `search` tool with queries related to the research topic. It returns relevant papers with complete metadata when available, including id, title, authors, organization, venue, year, n_citation, abstract, and a rule-based relevance score.
2. Use the `get_reference` tool on high-quality seed papers to retrieve their reference lists for backward-citation snowballing.
3. Use the `add_to_paper_set` tool at any time to add relevant papers to the collection. The collection is automatically deduplicated by paper id.
4. Make multiple rounds of tool calls to enrich the paper collection. Keep the total number of ordinary tool calls within about 20 rounds when possible.
5. The final target is more than 400 papers when the available tool results support it. Do not fabricate papers. If the tools cannot provide enough papers, return all collected papers.
6. You must terminate by calling `END` within 50 rounds. Do not enter an infinite tool-calling loop.

Tool descriptions:

* `search`: input `{"query": "search terms", "size": 20}`. It returns `[{"id": ..., "title": ..., "year": ..., "n_citation": ..., "score": ...}, ...]`. Choose `size` as needed; use no more than 20 papers per search.
* `get_reference`: input `{"aminer_ids": ["id", ...], "size_per_paper": 20}`. It retrieves reference lists for the given papers and returns `[{"id": ..., "title": ..., "year": ..., "n_citation": ..., "score": ...}, ...]`.
* `add_to_paper_set`: input `{"papers": ["id_1", "id_2", ...]}`. Add selected papers to the collection. Batch addition is supported. Prefer papers that are highly relevant to the topic, published in top journals or conferences, or highly cited.
* `END`: finish the process.

Output exactly one JSON object whenever you call a tool:

{"tool": "search", "params": {"query": "...", "size": 20}}
{"tool": "get_reference", "params": {"aminer_ids": ["id", "..."], "size_per_paper": 20}}
{"tool": "add_to_paper_set", "params": {"papers": ["id_1", "id_2"]}}
{"tool": "END"}

Query expansion and fallback strategies:

* If `search` returns fewer than 5 results or the average score is very low, the next step must be one of:
  1. Expand the query using synonyms, aliases, methods, datasets, benchmarks, or subfield terms related to the topic, and call `search` again.
  2. Call `get_reference` on the top 1-5 high-scoring papers from the latest search.
  3. If several searches fail, broaden the query by removing restrictive terms or using English aliases and abbreviations.
* For each automatic expansion, try at most two expanded queries before moving to reference expansion.

Search strategy:

* Start with broad and precise keyword searches for the topic.
* Add qualified seed papers to the paper set.
* After adding seed papers, prioritize `get_reference` on the strongest seeds to expand the collection.
* If search results are duplicate-heavy or mediocre, continue snowballing through references rather than repeating nearly identical searches.
* Only use information returned by the tools. Do not invent paper ids, titles, or metadata.
* Do not misinterpret the topic.
* Only one tool may be called at a time.
* If your previous output was not valid tool-call JSON, immediately correct it with a valid JSON object.
"""

from __future__ import annotations

import argparse
import json
import os
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from api_client import APIClient
from citation import get_reference_papers
from paper_set import PaperSet, paper_id_of
from prompt import REACT_SYSTEM_PROMPT
from search import search_papers


CURRENT_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = CURRENT_DIR / "outputs"


def default_llm_api_key() -> str | None:
    return os.getenv("LLM_API_KEY") or os.getenv("llm.api_key")


def default_llm_base_url() -> str | None:
    return os.getenv("LLM_BASE_URL") or os.getenv("llm.base_url")


def default_llm_model() -> str | None:
    return os.getenv("LLM_MODEL") or os.getenv("llm.model")


def extract_tool_call(text: str) -> dict[str, Any] | None:
    stripped = text.strip()
    if not stripped:
        return None

    candidates = [stripped]
    fenced = re.findall(r"```(?:json)?\s*(\{.*?\})\s*```", stripped, flags=re.DOTALL)
    candidates.extend(fenced)

    first = stripped.find("{")
    last = stripped.rfind("}")
    if first >= 0 and last > first:
        candidates.append(stripped[first : last + 1])

    for candidate in candidates:
        try:
            parsed = json.loads(candidate)
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict) and "tool" in parsed:
            return parsed
    return None


def is_tool_call_json(text: str) -> bool:
    return extract_tool_call(text) is not None


def compact_papers_for_model(papers: list[dict[str, Any]], *, limit: int = 40) -> list[dict[str, Any]]:
    compact: list[dict[str, Any]] = []
    for paper in papers[:limit]:
        compact.append(
            {
                "id": paper.get("id"),
                "title": paper.get("title"),
                "year": paper.get("year"),
                "n_citation": paper.get("n_citation"),
                "score": paper.get("score"),
            }
        )
    return compact


class ReactPaperCollector:
    def __init__(
        self,
        *,
        topic: str,
        api_key: str | None,
        base_url: str | None = None,
        models: list[str] | None = None,
        timeout: float = 300,
        max_rounds: int = 50,
        max_tool_calls: int = 20,
        target_size: int = 400,
        include_abstracts: bool = False,
    ) -> None:
        self.topic = topic.strip()
        if not self.topic:
            raise ValueError("Topic must be non-empty.")
        self.models = models
        self.max_rounds = max_rounds
        self.max_tool_calls = max_tool_calls
        self.target_size = target_size
        self.include_abstracts = include_abstracts
        self.client = APIClient(api_key=api_key, base_url=base_url, timeout=timeout)
        self.paper_set = PaperSet()
        self.paper_cache: dict[str, dict[str, Any]] = {}
        self.messages: list[dict[str, str]] = [
            {"role": "system", "content": REACT_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"Research topic: {self.topic}\n"
                    f"Current collection size: 0. Begin collection now."
                ),
            },
        ]
        self.tool_calls = 0

    def _remember_papers(self, papers: list[dict[str, Any]]) -> None:
        for paper in papers:
            paper_id = paper_id_of(paper).strip()
            if paper_id:
                self.paper_cache[paper_id] = paper

    def _tool_result_message(self, tool: str, payload: dict[str, Any]) -> dict[str, str]:
        return {
            "role": "user",
            "content": (
                f"Tool `{tool}` result:\n"
                f"{json.dumps(payload, ensure_ascii=False)}\n"
                "Continue with exactly one valid tool-call JSON object."
            ),
        }

    def execute_tool(self, call: dict[str, Any]) -> tuple[bool, dict[str, Any]]:
        tool = str(call.get("tool") or "").strip()
        params = call.get("params") if isinstance(call.get("params"), dict) else {}

        if tool == "search":
            query = str(params.get("query") or "").strip()
            size = int(params.get("size") or 20)
            papers = search_papers(query, size=size)
            self._remember_papers(papers)
            self.tool_calls += 1
            return False, {
                "query": query,
                "count": len(papers),
                "papers": compact_papers_for_model(papers),
            }

        if tool == "get_reference":
            ids = params.get("aminer_ids") or params.get("ids") or []
            if not isinstance(ids, list):
                ids = []
            size_per_paper = int(params.get("size_per_paper") or 20)
            papers = get_reference_papers(ids, topic=self.topic, size_per_paper=size_per_paper)
            self._remember_papers(papers)
            self.tool_calls += 1
            return False, {
                "seed_ids": ids,
                "count": len(papers),
                "papers": compact_papers_for_model(papers),
            }

        if tool == "add_to_paper_set":
            papers_or_ids = params.get("papers", call.get("papers", []))
            if not isinstance(papers_or_ids, list):
                papers_or_ids = []
            dict_papers = [item for item in papers_or_ids if isinstance(item, dict)]
            id_items = [str(item) for item in papers_or_ids if not isinstance(item, dict)]
            if dict_papers:
                self._remember_papers(dict_papers)
                self.paper_set.add_papers(dict_papers)
            if id_items:
                self.paper_set.add_ids(id_items, self.paper_cache)
            self.tool_calls += 1
            all_papers = self.paper_set.all_papers()
            return False, {
                "collection_size": len(self.paper_set),
                "papers": compact_papers_for_model(all_papers),
                "target_size": self.target_size,
            }

        if tool == "END":
            return True, {
                "collection_size": len(self.paper_set),
                "papers": self.paper_set.output(include_abstracts=self.include_abstracts),
            }

        return False, {
            "error": f"Unknown tool `{tool}`.",
            "valid_tools": ["search", "get_reference", "add_to_paper_set", "END"],
        }

    def run(self) -> list[dict[str, Any]]:
        for round_index in range(1, self.max_rounds + 1):
            if self.tool_calls >= self.max_tool_calls:
                print(
                    f"Ordinary tool-call budget ({self.max_tool_calls}) reached; "
                    "returning collected papers."
                )
                break

            response, ok = self.client.call_messages(
                self.messages,
                model_list=self.models,
                validator=is_tool_call_json,
            )
            if not ok:
                break

            self.messages.append({"role": "assistant", "content": response.content})
            call = extract_tool_call(response.content)
            if call is None:
                self.messages.append(
                    {
                        "role": "user",
                        "content": "Your output was not valid tool-call JSON. Return exactly one valid JSON object.",
                    }
                )
                continue

            done, result = self.execute_tool(call)
            print(f"Round {round_index}: tool={call.get('tool')} collection={len(self.paper_set)}")
            if done:
                return self.paper_set.output(include_abstracts=self.include_abstracts)
            self.messages.append(self._tool_result_message(str(call.get("tool")), result))

        return self.paper_set.output(include_abstracts=self.include_abstracts)

    def save_output(self, papers: list[dict[str, Any]]) -> Path:
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        safe_topic = re.sub(r"[^A-Za-z0-9._-]+", "_", self.topic).strip("_")[:80] or "topic"
        output_path = OUTPUT_DIR / f"{timestamp}_{safe_topic}.json"
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with output_path.open("w", encoding="utf-8") as file:
            json.dump(papers, file, ensure_ascii=False, indent=2)
        return output_path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Collect survey papers with an LLM-driven ReAct loop.")
    parser.add_argument("--topic", required=True, help="Research topic to collect papers for.")
    parser.add_argument(
        "--api-key",
        default=default_llm_api_key(),
        help="OpenAI-compatible LLM API key. Defaults to env LLM_API_KEY (legacy: llm.api_key).",
    )
    parser.add_argument(
        "--base-url",
        default=default_llm_base_url(),
        help="OpenAI-compatible LLM base URL. Defaults to env LLM_BASE_URL (legacy: llm.base_url).",
    )
    parser.add_argument("--models", nargs="*", default=None, help="Required model fallback list unless env LLM_MODEL (legacy: llm.model) is configured.")
    parser.add_argument("--timeout", type=float, default=300, help="Per-request model timeout in seconds.")
    parser.add_argument("--max-rounds", type=int, default=50)
    parser.add_argument("--max-tool-calls", type=int, default=20)
    parser.add_argument("--target-size", type=int, default=400)
    parser.add_argument("--include-abstracts", action="store_true")
    args = parser.parse_args()
    if args.models is None:
        model = default_llm_model()
        if model:
            args.models = [item.strip() for item in model.split(",") if item.strip()]
    return args


def main() -> None:
    args = parse_args()
    collector = ReactPaperCollector(
        topic=args.topic,
        api_key=args.api_key,
        base_url=args.base_url,
        models=args.models,
        timeout=args.timeout,
        max_rounds=args.max_rounds,
        max_tool_calls=args.max_tool_calls,
        target_size=args.target_size,
        include_abstracts=args.include_abstracts,
    )
    papers = collector.run()
    output_path = collector.save_output(papers)
    print(json.dumps(papers, ensure_ascii=False))
    print(f"\nSaved {len(papers)} papers to {output_path}")


if __name__ == "__main__":
    main()

openai>=1.68.0
requests>=2.31.0

from __future__ import annotations

import json
from typing import Any, Sequence

import requests

import _utils


AMINER_SEARCH_URL = "https://datacenter.aminer.cn/gateway/api/v3/paper/search/paper/SearchPro"


def _auth_headers() -> dict[str, str]:
    return {
        "Content-Type": "application/json;charset=utf-8",
        "Authorization": f"Bearer {_utils.get_aminer_key()}",
    }


def _extract_search_items(response_json: dict[str, Any]) -> list[dict[str, Any]]:
    data = response_json.get("data", [])
    if isinstance(data, dict):
        data = data.get("data") or data.get("items") or data.get("results") or []
    if not isinstance(data, list):
        return []
    return [item for item in data if isinstance(item, dict)]


def aminer_pro_search(
    query: str,
    use_topic: bool = True,
    year: int | None = None,
    size: int = 20,
    offset: int = 0,
) -> list[dict[str, Any]]:
    payload: dict[str, Any] = {
        "use_topic": use_topic,
        "query": query,
        "size": max(1, min(int(size), 100)),
        "offset": max(0, int(offset)),
        "end_year": int(year or 2026),
    }
    try:
        response = requests.post(
            AMINER_SEARCH_URL,
            headers=_auth_headers(),
            data=json.dumps(payload),
            timeout=(10, 30),
        )
        if response.status_code != 200:
            print(f"AMiner search failed: status={response.status_code}, detail={response.text[:300]}")
            return []
        return _extract_search_items(response.json())
    except (requests.RequestException, ValueError) as exc:
        print(f"AMiner search failed for query `{query}`: {exc}")
        return []


def search_papers(query: str, *, size: int = 20, year: int | None = None) -> list[dict[str, Any]]:
    size = max(1, min(int(size), 20))
    raw_items = aminer_pro_search(query, use_topic=True, year=year, size=size, offset=0)
    if not raw_items:
        return []

    ids = _utils.dedupe_preserve_order(_utils.extract_paper_id(item) for item in raw_items)
    details_by_id = {
        _utils.extract_paper_id(detail): detail
        for detail in _utils.aminer_get_paper_info_batch(ids)
        if _utils.extract_paper_id(detail)
    }

    papers: list[dict[str, Any]] = []
    for raw in raw_items:
        paper_id = _utils.extract_paper_id(raw)
        merged = dict(raw)
        if paper_id in details_by_id:
            merged.update(details_by_id[paper_id])
        normalized = _utils.normalize_paper_detail(merged, query=query)
        if normalized["id"] and normalized["title"]:
            papers.append(normalized)

    papers.sort(
        key=lambda item: (float(item.get("score", 0.0)), _utils.safe_int(item.get("n_citation"), 0)),
        reverse=True,
    )
    return papers[:size]


def search_adding(
    keyword_list: Sequence[str],
    topic: str,
    total_paper_details: Sequence[Any] | None = None,
    **_: Any,
) -> list[dict[str, Any]]:
    existing_ids = {
        _utils.extract_paper_id(item)
        for item in (total_paper_details or [])
        if _utils.extract_paper_id(item)
    }
    papers: list[dict[str, Any]] = []
    seen = set(existing_ids)
    for keyword in keyword_list:
        for paper in search_papers(str(keyword), size=20):
            paper_id = _utils.extract_paper_id(paper)
            if not paper_id or paper_id in seen:
                continue
            paper["topic"] = topic
            seen.add(paper_id)
            papers.append(paper)
    return papers


keywords_adding = search_adding


__all__ = [
    "aminer_pro_search",
    "keywords_adding",
    "search_adding",
    "search_papers",
]