yu-zhi-ran/scripts/opencode_search.py

#!/usr/bin/env python3
"""
通过 opencode CLI 执行联网搜索
利用 opencode 的 webfetch 能力（当前 AI 环境可无障碍访问互联网）

用法:
  python3 scripts/opencode_search.py --query "可持续生活 趋势 2026"
  python3 scripts/opencode_search.py --refresh-cache   # 刷新所有分类的缓存
"""
import argparse, datetime, json, logging, os, re, subprocess, sys, time
from pathlib import Path
from typing import Dict, List, Optional

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

LOGS_DIR = PROJECT_ROOT / "automation" / "logs"
TODAY = datetime.datetime.now().strftime("%Y-%m-%d")
LOG_FILE = LOGS_DIR / f"opencode_search_{TODAY}.log"

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler(LOG_FILE, encoding='utf-8'), logging.StreamHandler()])
logger = logging.getLogger(__name__)

SEARCH_CACHE_FILE = PROJECT_ROOT / "automation" / "data" / "search_cache.json"
SESSION_FILE = PROJECT_ROOT / "automation" / "data" / "opencode_session.txt"


def _get_or_create_session() -> Optional[str]:
    """获取或创建持久 session ID"""
    if SESSION_FILE.exists():
        try:
            sid = SESSION_FILE.read_text().strip()
            if sid:
                result = subprocess.run(
                    ["npx", "opencode", "run", "ping", "--session", sid, "--format", "json"],
                    capture_output=True, text=True, timeout=10,
                    cwd=str(PROJECT_ROOT),
                    env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
                )
                if result.returncode == 0:
                    return sid
        except Exception:
            pass
    result = subprocess.run(
        ["npx", "opencode", "run", "init", "--format", "json"],
        capture_output=True, text=True, timeout=30,
        cwd=str(PROJECT_ROOT),
        env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
    )
    for line in result.stdout.strip().split("\n"):
        try:
            event = json.loads(line)
            sid = event.get("sessionID") or event.get("part", {}).get("sessionID")
            if sid:
                SESSION_FILE.write_text(sid)
                return sid
        except Exception:
            pass
    return None


_session_id = None


def _run_opencode(prompt: str, timeout: int = 60) -> Optional[str]:
    """调用 opencode run 执行任务，返回文本输出"""
    global _session_id
    if _session_id is None:
        _session_id = _get_or_create_session()
    args = ["npx", "opencode", "run", prompt, "--format", "json"]
    if _session_id:
        args.extend(["--session", _session_id, "--continue"])
    try:
        result = subprocess.run(
            args,
            capture_output=True, text=True, timeout=timeout,
            cwd=str(PROJECT_ROOT),
            env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
        )
        if result.returncode != 0:
            logger.warning(f"opencode run 返回非零: {result.stderr[:200]}")
            return None
        for line in result.stdout.strip().split("\n"):
            try:
                event = json.loads(line)
                if event.get("type") == "error":
                    logger.warning(f"opencode 错误: {event}")
                    return None
            except json.JSONDecodeError:
                pass
        lines = []
        for line in result.stdout.strip().split("\n"):
            try:
                event = json.loads(line)
                if event.get("type") == "text":
                    text = event.get("part", {}).get("text", "")
                    if text:
                        lines.append(text)
            except json.JSONDecodeError:
                pass
        output = "\n".join(lines).strip()
        return output if output else None
    except subprocess.TimeoutExpired:
        logger.warning(f"opencode run 超时 ({timeout}s)")
        return None
    except Exception as e:
        logger.warning(f"opencode run 失败: {e}")
        return None


def search_via_opencode(query: str, max_results: int = 5) -> List[Dict]:
    """通过 MCP 搜索工具联网搜索（替代脆弱的 npx prompt 方式）"""
    try:
        from search_utils import search
        return search(query, max_results)
    except Exception as e:
        logger.warning("search_utils 不可用，回退子进程: %s", e)
        result = subprocess.run(
            [sys.executable, str(PROJECT_ROOT / "scripts" / "mcp_search_server.py"),
             "--query", query],
            capture_output=True, text=True, timeout=90,
        )
        if result.returncode == 0:
            try:
                return json.loads(result.stdout)[:max_results]
            except Exception:
                pass
    return []


def refresh_cache():
    """刷新所有搜索分类的缓存"""
    try:
        with open(PROJECT_ROOT / "config" / "sources.yaml") as f:
            import yaml
            cfg = yaml.safe_load(f)
        queries = [s["query"] for s in cfg["sustainability_sources"]["web_search"]]
    except Exception:
        logger.warning("无法读取 sources.yaml，使用默认查询")
        queries = [
            "以旧换新 二手交易 循环 2026",
            "新能源车 绿色通勤 低碳 2026",
            "干净饮食 有机食品 2026",
            "零浪费 极简生活 可持续时尚 2026",
            "绿色家电 一级能效 节能 2026",
            "碳账户 碳普惠 个人碳减排 2026",
            "环保科技 绿色产品 可持续材料 2026",
            "AI工具 人工智能 效率提升 2026",
        ]

    cache = {"_metadata": {"updated_at": datetime.datetime.now().isoformat()}}
    if SEARCH_CACHE_FILE.exists():
        try:
            old = json.loads(SEARCH_CACHE_FILE.read_text(encoding="utf-8"))
            for k, v in old.items():
                if not k.startswith("_"):
                    cache.setdefault(k, v)
        except Exception:
            pass

    for i, q in enumerate(queries):
        logger.info(f"[{i+1}/{len(queries)}] 搜索: {q}")
        results = search_via_opencode(q, max_results=4)
        if results:
            cache[q] = results
        else:
            logger.warning(f"  {q} 搜索无结果，保留旧缓存")
        time.sleep(2)

    SEARCH_CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
    SEARCH_CACHE_FILE.write_text(json.dumps(cache, ensure_ascii=False, indent=2), encoding="utf-8")
    logger.info(f"缓存已刷新: {sum(len(v) for v in cache.values())} 条")


def main():
    parser = argparse.ArgumentParser(description="通过 opencode 联网搜索")
    parser.add_argument("--query", help="搜索词")
    parser.add_argument("--refresh-cache", action="store_true", help="刷新所有分类缓存")
    parser.add_argument("--max-results", type=int, default=5)
    args = parser.parse_args()

    if args.refresh_cache:
        refresh_cache()
        return

    if args.query:
        results = search_via_opencode(args.query, args.max_results)
        print(json.dumps(results, ensure_ascii=False, indent=2))
        return

    parser.print_help()


if __name__ == "__main__":
    main()