Agent/cyber_agent/tools/builtin/memory.py

"""
Memory 相关工具 —— 目前只包含 dream 操作（见 cyber_agent/docs/memory.md 第四节）。

dream 整理 Agent 身份的长期记忆：回顾最近 trace 的执行历史，
逐个 trace 做反思，再跨 trace 整合写回记忆文件。

设计要点：
- 需要 config.memory（MemoryConfig）才可用；否则报错。
- 不是 knowledge_save_pending 那样每 trace 都要用的日常工具 ——
  所以放在独立 group "memory"，通过 tool_groups 显式开启。
"""

from __future__ import annotations

import logging
from typing import Optional

from cyber_agent.core.dream import DreamScope, run_dream
from cyber_agent.core.memory import MemoryConfig
from cyber_agent.tools import tool, ToolResult, ToolContext

logger = logging.getLogger(__name__)


@tool(groups=["memory"], hidden_params=["context"])
async def dream(
    reflect_model: str = "",
    dream_model: str = "",
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """整理长期记忆。回顾最近的执行历史，更新记忆文件。

    本工具做两件事：
        1. per-trace 反思：扫描未反思的 trace，为每个生成反思摘要
        2. 跨 trace 整合：汇总未消化的反思 + 当前记忆，让 LLM 更新记忆文件

    需要 RunConfig.memory（MemoryConfig）才可调用。

    Args:
        reflect_model: per-trace 反思用的模型（空则默认 gpt-4o-mini）
        dream_model:   跨 trace 整合用的模型（空则默认 gpt-4o）
    """
    runner = context.get("runner") if context else None
    if runner is None:
        return ToolResult(
            title="❌ dream 不可用",
            output="缺少 runner（需要从 AgentRunner 上下文调用）",
            error="runner not in context",
        )

    memory_config = context.get("memory_config") if context else None
    if not isinstance(memory_config, MemoryConfig):
        return ToolResult(
            title="❌ dream 不可用",
            output="当前工具调用未注入 MemoryConfig，不是 memory-bearing Agent",
            error="memory_config not in tool context",
        )

    dream_scope = context.get("dream_scope") if context else None
    if not isinstance(dream_scope, DreamScope):
        return ToolResult(
            title="❌ dream 不可用",
            output="当前工具调用未注入 DreamScope，为避免跨身份扫描已拒绝执行",
            error="dream_scope not in tool context",
        )

    if not runner.trace_store or not runner.llm_call:
        return ToolResult(
            title="❌ dream 不可用",
            output="runner 缺少 trace_store 或 llm_call",
            error="runner dependencies missing",
        )

    report = await run_dream(
        store=runner.trace_store,
        llm_call=runner.llm_call,
        memory_config=memory_config,
        dream_scope=dream_scope,
        reflect_model=reflect_model or "gpt-4o-mini",
        dream_model=dream_model or "gpt-4o",
    )

    lines = []
    lines.append(f"per-trace 反思: {len(report.per_trace_summaries)} 条")
    if report.skipped_traces:
        lines.append(f"跳过: {len(report.skipped_traces)} 条 trace（日志详见 logger）")
    lines.append(f"消化 reflection: {report.consumed_reflection_count} 条")
    lines.append(f"更新记忆文件: {len(report.updated_files)} 个")
    for p in report.updated_files:
        lines.append(f"  - {p}")
    if report.reasoning:
        lines.append(f"\n整合理由: {report.reasoning}")

    output = "\n".join(lines)
    return ToolResult(
        title="🧠 dream 完成",
        output=output,
        long_term_memory=f"dream: reflected={len(report.per_trace_summaries)}, "
                         f"consumed={report.consumed_reflection_count}, "
                         f"files_updated={len(report.updated_files)}",
        metadata={
            "per_trace_count": len(report.per_trace_summaries),
            "consumed": report.consumed_reflection_count,
            "updated_files": report.updated_files,
        },
    )