Video-Tree-TRM5/.claude/tools/semantic_scholar_fetch.py

"""Semantic Scholar 搜索命令行工具。"""

from __future__ import annotations

import argparse
import json
import sys
import time
from typing import Any

try:
    import requests
except ImportError:
    print("警告: 未安装 requests，已跳过 Semantic Scholar CLI。", file=sys.stderr)
    sys.exit(0)


API_URL = "https://api.semanticscholar.org/graph/v1/paper/search"
REQUEST_FIELDS = "title,authors,year,venue,citationCount,externalIds,tldr"
REQUEST_INTERVAL_SECONDS = 1.0
REQUEST_TIMEOUT_SECONDS = 30.0
USER_AGENT = "semantic_scholar_fetch.py/1.0"
_LAST_REQUEST_AT: float | None = None


def _normalize_filter_values(values: list[str] | None) -> str | None:
    """将过滤参数整理为 API 需要的逗号分隔字符串。"""
    if not values:
        return None

    items: list[str] = []
    for value in values:
        for item in value.split(","):
            cleaned = item.strip()
            if cleaned:
                items.append(cleaned)

    if not items:
        return None
    return ",".join(items)


def _respect_rate_limit() -> None:
    """确保连续请求之间至少间隔一秒。"""
    global _LAST_REQUEST_AT

    now = time.monotonic()
    if _LAST_REQUEST_AT is not None:
        elapsed = now - _LAST_REQUEST_AT
        if elapsed < REQUEST_INTERVAL_SECONDS:
            time.sleep(REQUEST_INTERVAL_SECONDS - elapsed)

    _LAST_REQUEST_AT = time.monotonic()


def _clean_text(value: Any) -> str:
    """压缩并清理任意文本值。"""
    return " ".join(str(value).split())


def _normalize_authors(authors: Any) -> list[str]:
    """将作者字段统一为作者姓名列表。"""
    normalized_authors: list[str] = []
    if not isinstance(authors, list):
        return normalized_authors

    for author in authors:
        if isinstance(author, dict):
            name = author.get("name")
            if name:
                normalized_authors.append(_clean_text(name))
        elif author:
            normalized_authors.append(_clean_text(author))

    return normalized_authors


def _normalize_external_ids(external_ids: Any) -> dict[str, str]:
    """将 externalIds 统一整理为字符串字典，并保留 arXiv 标识。"""
    normalized: dict[str, str] = {}
    if isinstance(external_ids, dict):
        for key, value in external_ids.items():
            if value is None:
                continue
            cleaned_value = _clean_text(value)
            if cleaned_value:
                normalized[str(key)] = cleaned_value

    arxiv_id = (
        normalized.get("ArXiv") or normalized.get("arXiv") or normalized.get("arxiv")
    )
    if arxiv_id and "ArXiv" not in normalized:
        normalized["ArXiv"] = arxiv_id

    return normalized


def _normalize_tldr(tldr: Any) -> Any:
    """保留 TLDR 字段原始结构，但去除明显的空字符串。"""
    if isinstance(tldr, dict):
        normalized_tldr: dict[str, Any] = {}
        for key, value in tldr.items():
            if isinstance(value, str):
                cleaned = _clean_text(value)
                if cleaned:
                    normalized_tldr[key] = cleaned
            elif value is not None:
                normalized_tldr[key] = value
        return normalized_tldr or None
    return tldr


def _normalize_paper(paper: dict[str, Any]) -> dict[str, Any]:
    """将 API 返回的论文记录压缩成稳定的输出结构。"""
    return {
        "title": _clean_text(paper.get("title", "")),
        "authors": _normalize_authors(paper.get("authors")),
        "year": paper.get("year"),
        "venue": _clean_text(paper.get("venue", "")),
        "citationCount": paper.get("citationCount"),
        "externalIds": _normalize_external_ids(paper.get("externalIds")),
        "tldr": _normalize_tldr(paper.get("tldr")),
    }


def _request_json(url: str, params: dict[str, Any]) -> dict[str, Any]:
    """发起受限频率控制的 GET 请求并返回 JSON 对象。"""
    _respect_rate_limit()
    response = requests.get(
        url,
        params=params,
        headers={"User-Agent": USER_AGENT},
        timeout=REQUEST_TIMEOUT_SECONDS,
    )
    response.raise_for_status()

    payload = response.json()
    if not isinstance(payload, dict):
        raise ValueError("Semantic Scholar 返回了非对象类型的 JSON")
    return payload


def search_papers(
    query: str,
    max_results: int,
    fields_of_study: list[str] | None = None,
    publication_types: list[str] | None = None,
) -> list[dict[str, Any]]:
    """搜索 Semantic Scholar 论文并返回标准化结果。"""
    params: dict[str, Any] = {
        "query": query,
        "limit": max_results,
        "fields": REQUEST_FIELDS,
    }

    normalized_fields_of_study = _normalize_filter_values(fields_of_study)
    if normalized_fields_of_study is not None:
        params["fieldsOfStudy"] = normalized_fields_of_study

    normalized_publication_types = _normalize_filter_values(publication_types)
    if normalized_publication_types is not None:
        params["publicationTypes"] = normalized_publication_types

    payload = _request_json(API_URL, params)
    papers = payload.get("data", [])
    if not isinstance(papers, list):
        raise ValueError("Semantic Scholar 响应缺少 data 列表")

    normalized_papers: list[dict[str, Any]] = []
    for paper in papers:
        if isinstance(paper, dict):
            normalized_papers.append(_normalize_paper(paper))
    return normalized_papers


def build_parser() -> argparse.ArgumentParser:
    """构建命令行参数解析器。"""
    parser = argparse.ArgumentParser(description="Semantic Scholar 搜索工具")
    subparsers = parser.add_subparsers(dest="command", required=True)

    search_parser = subparsers.add_parser("search", help="搜索论文")
    search_parser.add_argument("query", help="搜索关键词")
    search_parser.add_argument("--max", dest="max_results", type=int, default=10)
    search_parser.add_argument(
        "--fields-of-study",
        nargs="+",
        default=None,
        help="按学科领域过滤",
    )
    search_parser.add_argument(
        "--publication-types",
        nargs="+",
        default=None,
        help="按出版类型过滤",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    """CLI 入口。"""
    parser = build_parser()
    args = parser.parse_args(argv)

    try:
        if args.command == "search":
            results = search_papers(
                query=args.query,
                max_results=args.max_results,
                fields_of_study=args.fields_of_study,
                publication_types=args.publication_types,
            )
            json.dump(results, sys.stdout, ensure_ascii=False, indent=2)
            sys.stdout.write("\n")
            return 0
    except (requests.RequestException, ValueError) as exc:
        print(f"错误: {exc}", file=sys.stderr)
        return 1

    parser.error(f"未知命令: {args.command}")
    return 2


if __name__ == "__main__":
    sys.exit(main())