aiagent/backend/app/core/streamlined_output.py

"""
工具结果流式美化引擎 — Streamlined Output

参考 Claude Code:
- src/utils/streamlinedTransform.ts  — 累积计数 + 文本断点
- src/utils/collapseReadSearch.ts    — 搜索/读取折叠分组
- src/utils/groupToolUses.ts         — 同类型工具归并

将工具调用过程转译为自然语言描述，让用户看到简洁摘要而非原始 JSON。
"""
from __future__ import annotations

import logging
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set, Tuple

logger = logging.getLogger(__name__)


# ──────────────────────────── 工具分类 ────────────────────────────

class ToolCategory(str, Enum):
    SEARCH = "search"    # 搜索类 (grep, glob, web_search)
    READ = "read"        # 读取类 (file_read, list_files)
    WRITE = "write"      # 写入类 (file_write, file_edit, deploy)
    COMMAND = "command"  # 执行类 (bash, code_execute, database)
    OTHER = "other"      # 其他


# 参考 Claude Code streamlinedTransform.ts L38-50
SEARCH_TOOLS: Set[str] = {
    "grep", "search_content", "search_files", "search_code",
    "web_search", "web_fetch", "find_files", "rg", "rg_search",
}

READ_TOOLS: Set[str] = {
    "file_read", "list_files", "read_file", "read_dir",
    "system_info", "entity_search", "knowledge_graph_search",
    "browser_use", "image_ocr", "image_vision",
    "url_parse", "http_request",
}

WRITE_TOOLS: Set[str] = {
    "file_write", "file_edit", "write_file", "edit_file",
    "deploy_push", "docker_manage", "notebook_edit",
    "pdf_generate", "excel_process",
}

COMMAND_TOOLS: Set[str] = {
    "code_execute", "database_query", "git_operation",
    "agent_create", "agent_call", "create_task", "task_plan",
    "adb_log", "crypto_util", "schedule_create", "schedule_delete",
    "project_scaffold", "tool_register",
    "feishu_create_doc", "feishu_create_sheet", "feishu_send_approval",
    "feishu_upload_file", "send_email",
    "bash", "shell", "cmd",
}


def categorize_tool(tool_name: str) -> ToolCategory:
    """根据工具名称分类。"""
    if not tool_name:
        return ToolCategory.OTHER
    name = tool_name.lower().strip()
    if name in SEARCH_TOOLS:
        return ToolCategory.SEARCH
    if name in READ_TOOLS:
        return ToolCategory.READ
    if name in WRITE_TOOLS:
        return ToolCategory.WRITE
    if name in COMMAND_TOOLS:
        return ToolCategory.COMMAND
    return ToolCategory.OTHER


# ──────────────────────────── 计数器 ────────────────────────────

class ToolCounts:
    """累积工具调用计数。"""
    __slots__ = ("searches", "reads", "writes", "commands", "other")

    def __init__(self):
        self.searches: int = 0
        self.reads: int = 0
        self.writes: int = 0
        self.commands: int = 0
        self.other: int = 0

    def add(self, category: ToolCategory) -> None:
        if category == ToolCategory.SEARCH:
            self.searches += 1
        elif category == ToolCategory.READ:
            self.reads += 1
        elif category == ToolCategory.WRITE:
            self.writes += 1
        elif category == ToolCategory.COMMAND:
            self.commands += 1
        else:
            self.other += 1

    def has_any(self) -> bool:
        return any((self.searches, self.reads, self.writes, self.commands, self.other))

    def reset(self) -> None:
        self.searches = 0
        self.reads = 0
        self.writes = 0
        self.commands = 0
        self.other = 0


# ──────────────────────────── 摘要生成 ────────────────────────────

def _plural_en(n: int, singular: str, plural: str = "") -> str:
    """英文复数（用于工具名）。"""
    if n == 1:
        return singular
    return plural or singular + "s"


def get_tool_summary_text(counts: ToolCounts) -> Optional[str]:
    """生成工具调用累计摘要文本（中文）。

    参考 Claude Code streamlinedTransform.ts L73-104
    """
    parts: List[str] = []

    if counts.searches > 0:
        parts.append(f"搜索了 {counts.searches} 个模式")
    if counts.reads > 0:
        parts.append(f"读取了 {counts.reads} 个文件")
    if counts.writes > 0:
        parts.append(f"写入了 {counts.writes} 个文件")
    if counts.commands > 0:
        parts.append(f"执行了 {counts.commands} 条命令")
    if counts.other > 0:
        parts.append(f"调用了 {counts.other} 个工具")

    if not parts:
        return None

    return "、".join(parts)


def get_search_read_summary(
    search_count: int,
    read_count: int,
    is_active: bool = False,
    list_count: int = 0,
) -> str:
    """生成搜索/读取操作的摘要文本（用于折叠组）。

    参考 Claude Code collapseReadSearch.ts L961-1066
    """
    parts: List[str] = []

    if search_count > 0:
        if is_active:
            verb = "正在搜索" if len(parts) == 0 else "搜索"
        else:
            verb = "已搜索" if len(parts) == 0 else "搜索了"
        parts.append(f"{verb} {search_count} 个模式")

    if read_count > 0:
        if is_active:
            verb = "正在读取" if len(parts) == 0 else "读取"
        else:
            verb = "已读取" if len(parts) == 0 else "读取了"
        parts.append(f"{verb} {read_count} 个文件")

    if list_count > 0:
        if is_active:
            verb = "正在列出" if len(parts) == 0 else "列出"
        else:
            verb = "已列出" if len(parts) == 0 else "列出了"
        parts.append(f"{verb} {list_count} 个目录")

    text = "，".join(parts)
    if is_active:
        text += "…"
    return text


# ──────────────────────────── 流式转换器 ────────────────────────────

class StreamlinedTransformer:
    """有状态的流式转换器：在文本消息之间累积工具计数。

    参考 Claude Code streamlinedTransform.ts L130-193

    用法:
        transformer = StreamlinedTransformer()
        for event in stream:
            transformed = transformer.transform(event)
            if transformed:
                yield transformed
    """

    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self._counts = ToolCounts()
        self._pending_tool_results: List[Dict[str, Any]] = []

    def transform(self, event: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """转换单个 SSE 事件。返回 None 表示该事件应被过滤。"""
        if not self.enabled:
            return event

        event_type = event.get("type", "")

        # ── 文本消息：直接输出，重置计数 ──
        if event_type in ("message", "final"):
            self._counts.reset()
            self._pending_tool_results.clear()
            return event

        # ── 思考/推理：保留（可配置是否过滤） ──
        if event_type == "think":
            # 保留 think 事件供前端展示推理过程
            return event

        # ── 工具调用：累积计数 ──
        if event_type == "tool_call":
            tool_name = event.get("name", "") or event.get("tool_name", "")
            category = categorize_tool(tool_name)
            self._counts.add(category)
            # 转发工具调用事件（带有分类信息）
            return {
                **event,
                "tool_category": category.value,
            }

        # ── 工具结果：暂存，等下次文本消息时清除 ──
        if event_type == "tool_result":
            self._pending_tool_results.append(event)
            # 发出累计摘要而非原始结果
            summary = get_tool_summary_text(self._counts)
            if summary:
                return {
                    "type": "streamlined_summary",
                    "summary": summary,
                    "counts": {
                        "searches": self._counts.searches,
                        "reads": self._counts.reads,
                        "writes": self._counts.writes,
                        "commands": self._counts.commands,
                        "other": self._counts.other,
                    },
                }
            return None

        # ── 错误 ──
        if event_type == "error":
            return event

        # ── 其他事件类型：保留 ──
        return event

    def flush(self) -> Optional[Dict[str, Any]]:
        """刷新最终的累计摘要。"""
        summary = get_tool_summary_text(self._counts)
        if summary:
            return {
                "type": "streamlined_summary",
                "summary": summary,
                "counts": {
                    "searches": self._counts.searches,
                    "reads": self._counts.reads,
                    "writes": self._counts.writes,
                    "commands": self._counts.commands,
                    "other": self._counts.other,
                },
            }
        return None

    def reset(self) -> None:
        """重置所有累积状态。"""
        self._counts.reset()
        self._pending_tool_results.clear()


# ──────────────────────────── 批量折叠器 ────────────────────────────

class ReadSearchCollapser:
    """折叠连续的搜索/读取操作为单个摘要组。

    参考 Claude Code collapseReadSearch.ts — 按消息流折叠
    连续 search/read 类型的工具使用和结果。

    规则:
    - 连续的 search/read 工具调用会合并为一个组
    - writer/command 工具调用会打断组
    - 文本消息会打断组
    """

    def __init__(self):
        self._group: List[Dict[str, Any]] = []
        self._search_count = 0
        self._read_count = 0
        self._list_count = 0
        self._is_active = False

    def feed_tool_call(self, tool_name: str, tool_input: Optional[Dict] = None) -> Optional[Dict[str, Any]]:
        """处理一个工具调用，如果组被打断则返回 flushed 的组摘要。"""
        cat = categorize_tool(tool_name)

        if cat in (ToolCategory.SEARCH, ToolCategory.READ):
            # 可折叠：加入当前组
            if cat == ToolCategory.SEARCH:
                self._search_count += 1
            else:
                self._read_count += 1
            self._group.append({"name": tool_name, "input": tool_input})
            self._is_active = True
            # 返回进行中的摘要
            return {
                "type": "collapsed_progress",
                "summary": get_search_read_summary(
                    self._search_count, self._read_count,
                    is_active=True, list_count=self._list_count,
                ),
                "search_count": self._search_count,
                "read_count": self._read_count,
                "list_count": self._list_count,
            }
        else:
            # 不可折叠：先 flush 组，再返回 None（调用方自行处理）
            flushed = self.flush()
            self._group = []
            self._is_active = False
            return flushed

    def flush(self) -> Optional[Dict[str, Any]]:
        """输出当前折叠组的最终摘要。"""
        if self._search_count == 0 and self._read_count == 0 and self._list_count == 0:
            return None
        result = {
            "type": "collapsed_group",
            "summary": get_search_read_summary(
                self._search_count, self._read_count,
                is_active=False, list_count=self._list_count,
            ),
            "search_count": self._search_count,
            "read_count": self._read_count,
            "list_count": self._list_count,
            "tool_count": len(self._group),
        }
        self._search_count = 0
        self._read_count = 0
        self._list_count = 0
        self._group = []
        self._is_active = False
        return result

    def reset(self) -> None:
        self._group = []
        self._search_count = 0
        self._read_count = 0
        self._list_count = 0
        self._is_active = False


# ──────────────────────────── 工厂函数 ────────────────────────────

def create_streamlined_transformer(enabled: bool = True) -> StreamlinedTransformer:
    """创建 StreamlinedTransformer 实例。"""
    return StreamlinedTransformer(enabled=enabled)


def create_read_search_collapser() -> ReadSearchCollapser:
    """创建 ReadSearchCollapser 实例。"""
    return ReadSearchCollapser()