myquant/finance/agents/orchestrator.py

"""
AgentOrchestrator — Agent 编排器。

管理所有 Agent 的生命周期、执行顺序、结果传递。
"""

from datetime import datetime

from agents.research_agent import ResearchAgent
from agents.selection_agent import SelectionAgent
from agents.risk_agent import RiskAgent
from agents.report_agent import ReportAgent


class AgentOrchestrator:
    """
    Agent 编排器。

    用法:
        orch = AgentOrchestrator(dm=dm, fe=engine_fe, bt=engine_bt, ...)
        orch.setup()  # 注册所有 Agent
        orch.run_daily()  # 执行每日流程
    """

    def __init__(self, **engines):
        self.engines = engines
        self.agents: dict = {}
        self._last_results: dict = {}

    def setup(self):
        """注册所有 Agent。"""
        self.agents["research"] = ResearchAgent(**self.engines)
        self.agents["selection"] = SelectionAgent(**self.engines)
        self.agents["risk"] = RiskAgent(**self.engines)
        self.agents["report"] = ReportAgent(**self.engines)
        print(f"[Orchestrator] 已注册 {len(self.agents)} 个 Agent: {list(self.agents)}")

    # ── 每日流程 ──────────────────────────────────────────

    def run_daily(self, date: str | None = None) -> dict:
        """
        每日任务流:

        1. 同步数据
        2. 风险评估
        3. 股票打分
        4. 生成日报
        """
        date = date or datetime.now().strftime("%Y%m%d")
        print(f"\n{'='*60}")
        print(f"[Orchestrator] 每日流程 — {date}")
        print(f"{'='*60}")

        results = {"date": date}

        # Step 1: 增量同步已缓存股票的最新行情
        print("\n[Step 1/4] 同步行情...")
        dm = self.engines.get("dm")
        sent = self.engines.get("sent")
        from database.dao import get_latest_trade_date

        scope_stocks = []
        if sent:
            scope_stocks = sent.get_scope_stocks()
        if not scope_stocks and dm:
            scope_stocks = list(dm.get_stock_list().index[:100])
        print("  范围: {} 只股票".format(len(scope_stocks)))

        # 分类：已缓存（增量更新），未缓存（统计跳过）
        cached = [c for c in scope_stocks if get_latest_trade_date(c)]
        uncached = len(scope_stocks) - len(cached)
        print("  已缓存: {} 只 (增量更新), 未缓存: {} 只 (跳过，需首次批量预热)".format(len(cached), uncached))

        synced = 0
        for i, ts_code in enumerate(cached):
            try:
                n = dm.sync_daily(ts_code)
                synced += n
            except Exception:
                continue
            if (i + 1) % 100 == 0:
                print("  [同步] 进度: {}/{}".format(i + 1, len(cached)))
        print("  同步完成: {} 条新数据 (已缓存{}/全量{})".format(synced, len(cached), len(scope_stocks)))

        if uncached > 0:
            print("  提示: {} 只股票未缓存，运行 'agent_cli.py warmup' 首次批量预热".format(uncached))

        # Step 2: 风险评估
        print("\n[Step 2/4] 风险评估...")
        risk = self.agents["risk"].execute()
        results["risk"] = risk
        print(f"  风险: {risk['risk_level']}, 仓位: {risk['target_exposure']:.0%}")

        # Step 3: 选股打分
        print("\n[Step 3/4] 股票打分...")
        selection = self.agents["selection"].execute(date=date, top_n=15)
        results["selection"] = selection
        top = selection.get("top_picks", [])
        if top:
            print("  Top 5: {}".format(", ".join(p['ts_code'] for p in top[:5])))

        # Step 4: 情绪因子
        print("\n[Step 4/5] 情绪因子...")
        sentiment_df = None
        sent_eng = self.engines.get("sent")
        if sent_eng:
            try:
                # 用范围内第一只有缓存的股票计算情绪因子
                ref_code = cached[0] if cached else "000001.SZ"
                sentiment_df = sent_eng.compute(ref_code, max_news=30)
                if sentiment_df is not None and not sentiment_df.empty:
                    valid = sentiment_df.dropna(how="all")
                    print("  {}: {} 个因子, {} 个有效交易日".format(ref_code, sentiment_df.shape[1], len(valid)))
                else:
                    print("  (无有效情绪数据)")
            except Exception as e:
                print("  [SKIP] 情绪因子计算失败: {}".format(e))
        else:
            print("  (SentimentEngine 未配置)")
        results["sentiment"] = sentiment_df

        # Step 5: 生成日报
        print("\n[Step 5/5] 生成日报...")
        report = self.agents["report"].execute(
            date=date,
            selection_result=selection,
            risk_result=risk,
            sentiment_result=sentiment_df,
        )
        results["report"] = report
        print(f"  日报: {report['report_path']}")

        self._last_results = results
        print(f"\n{'='*60}")
        print(f"[Orchestrator] 每日流程完成")
        print(f"{'='*60}\n")
        return results

    # ── 研究流程（每周一次） ────────────────────────────────

    def run_research_cycle(self, ts_codes: list[str] | None = None) -> dict:
        """
        研究周期:

        1. 因子发现与评估
        2. 更新 IC 权重
        """
        print(f"\n{'='*60}")
        print(f"[Orchestrator] 研究周期")
        print(f"{'='*60}")

        print("\n[Step 1/2] 因子发现...")
        research = self.agents["research"].execute(ts_codes=ts_codes)
        results = {"research": research}

        top = research.get("top_factors", [])
        if top:
            print(f"  Top 5 因子:")
            for f in top[:5]:
                print(f"    {f['name']:20s} IC={f['ic_mean']:+.4f}  ICIR={f['icir']:.3f}")

        return results

    # ── 便捷方法 ──────────────────────────────────────────

    def picks(self, date: str | None = None, top_n: int = 15) -> dict:
        """快速选股。"""
        return self.agents["selection"].execute(date=date, top_n=top_n)

    def risk_check(self) -> dict:
        """快速风险评估。"""
        return self.agents["risk"].execute()

    def generate_report(self, date: str | None = None) -> dict:
        """快速生成日报。"""
        date = date or datetime.now().strftime("%Y%m%d")

        # 尝试拉取最新数据（Tushare 优先，几秒即可完成）
        dm = self.engines.get("dm")
        if dm:
            try:
                dm.sync_daily("000001.SZ")
            except Exception:
                pass

        # 查 DB 最新交易日
        data_freshness = None
        try:
            from database.dao import get_latest_trade_date
            data_freshness = get_latest_trade_date("000001.SZ")
        except Exception:
            pass

        sel = self.picks(date)
        risk = self.risk_check()
        # 尝试取情绪因子
        sentiment_df = None
        sent_eng = self.engines.get("sent")
        if sent_eng:
            try:
                sentiment_df = sent_eng.compute("000001.SZ", max_news=30)
            except Exception:
                pass
        return self.agents["report"].execute(
            date=date, selection_result=sel, risk_result=risk,
            sentiment_result=sentiment_df, data_freshness=data_freshness,
        )

    @property
    def last_results(self) -> dict:
        return self._last_results