lambdagent/agentexample/agent67/core/assistant.py

"""
agent67.core.assistant — lambda 的大脑 🐑 (v2)

v2 升级：
  - 使用 lambdagent 内置 21 工具（替代手写 6 工具）
  - ToolGateway 安全网关（替代手写 safety.py）
  - Hook 系统全局审计
  - ContextManager 防状态爆炸
  - 保留 macOS 专属工具（browser, app, system, screenshot）

Lambda 表达式:
    lambda = Memory(Loop(Brain >> parse_action >> route_execute))
"""
from __future__ import annotations

import json
import re
import sys
import os
from pathlib import Path
from typing import Any

# 路径设置
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from lambdagent import Lam, Tool, Loop, Context
from lambdagent.extensions import Memory, Guard
from lambdagent.context_manager import ContextManager
from lambdagent.hooks import HookRegistry

# 内置工具
from lambdagent.builtin_tools.registry import BUILTIN_TOOLS, resolve_tools
from lambdagent.builtin_tools.terminal_ui import TerminalUI

from .claude_lam import ClaudeLam
from .ollama_lam import OllamaLam
from .config import BACKEND_OLLAMA, BACKEND_CLAUDE_CODE, BACKEND_API
from .prompt import SYSTEM_PROMPT

# macOS 专属工具（保留）
from ..tools.browser_controller import control_browser
from ..tools.app_controller import control_app
from ..tools.system_info import query_system
from ..tools.screenshot import take_screenshot


# ════════════════════════════════════════════════════════════
# 工具注册表 (v2: 内置工具 + macOS 工具)
# ════════════════════════════════════════════════════════════

def _build_tool_registry() -> dict:
    """合并内置工具 + macOS 专属工具"""
    # 内置 21 个工具
    tools = dict(BUILTIN_TOOLS)

    # macOS 专属（保留原有能力）
    tools["browser"] = Tool("browser", control_browser)
    tools["app"] = Tool("app", control_app)
    tools["system"] = Tool("system", query_system)
    tools["screenshot"] = Tool("screenshot", take_screenshot)

    # 完成信号
    tools["done"] = Tool("done", lambda x: x)

    return tools


TOOL_REGISTRY = _build_tool_registry()


# ════════════════════════════════════════════════════════════
# 解析 & 执行引擎 (v2)
# ════════════════════════════════════════════════════════════

def _extract_json_blocks(text: str) -> list[str]:
    """Extract JSON code blocks, handling nested ``` in content.

    Strategy: find ```json markers, then find the matching ``` closer
    by looking for a ``` that is followed by a newline or end-of-string
    (not inside a JSON string value).
    """
    blocks = []
    i = 0
    marker = "```json"
    while i < len(text):
        start = text.find(marker, i)
        if start == -1:
            break
        # Skip past the marker and optional newline
        content_start = start + len(marker)
        if content_start < len(text) and text[content_start] == '\n':
            content_start += 1

        # Find the closing ``` — look for ``` at start of line or after newline
        # that's NOT inside a JSON string
        best_end = -1
        j = content_start
        while j < len(text):
            pos = text.find("```", j)
            if pos == -1:
                break
            # Check if this ``` is a real closer (not inside JSON string content)
            # Heuristic: if the text between content_start and pos contains
            # a valid-looking JSON with "action", it's the closer
            candidate = text[content_start:pos].strip()
            if candidate and ('"action"' in candidate or '"tool"' in candidate):
                try:
                    # Try to parse — if it works, this is the right closer
                    _try_parse_json(candidate)
                    best_end = pos
                    break
                except Exception:
                    pass
                # Even if parse fails, if it looks like JSON, take it
                if candidate.startswith("{") and candidate.rstrip().endswith("}"):
                    best_end = pos
                    break
            j = pos + 3

        if best_end != -1:
            block = text[content_start:best_end].strip()
            if block:
                blocks.append(block)
            i = best_end + 3
        else:
            # Fallback: take everything until the next ```
            end = text.find("```", content_start)
            if end != -1:
                block = text[content_start:end].strip()
                if block:
                    blocks.append(block)
                i = end + 3
            else:
                break

    return blocks


def _try_parse_json(json_str: str) -> dict | None:
    """Try multiple strategies to parse potentially malformed JSON from LLM output."""
    # Strategy 1: direct parse
    try:
        return json.loads(json_str)
    except json.JSONDecodeError:
        pass

    # Strategy 2: fix newlines inside JSON string values
    # LLM outputs real newlines inside "content": "line1\nline2"
    # We need to escape them but ONLY inside string values
    try:
        fixed = _fix_json_newlines(json_str)
        return json.loads(fixed)
    except (json.JSONDecodeError, Exception):
        pass

    # Strategy 3: balanced brace extraction + newline fix
    try:
        depth = 0
        start = json_str.index("{")
        in_string = False
        escape = False
        for i in range(start, len(json_str)):
            c = json_str[i]
            if escape:
                escape = False
                continue
            if c == '\\':
                escape = True
                continue
            if c == '"':
                in_string = not in_string
                continue
            if not in_string:
                if c == '{':
                    depth += 1
                elif c == '}':
                    depth -= 1
                    if depth == 0:
                        candidate = json_str[start:i+1]
                        fixed = _fix_json_newlines(candidate)
                        return json.loads(fixed)
    except (json.JSONDecodeError, ValueError):
        pass

    return None


def _fix_json_newlines(s: str) -> str:
    """Replace real newlines inside JSON string values with \\n.

    Walks the string character by character, only replacing newlines
    when inside a quoted string value.
    """
    result = []
    in_string = False
    escape = False
    for c in s:
        if escape:
            result.append(c)
            escape = False
            continue
        if c == '\\':
            result.append(c)
            escape = True
            continue
        if c == '"':
            in_string = not in_string
            result.append(c)
            continue
        if c == '\n' and in_string:
            result.append('\\n')
            continue
        result.append(c)
    return ''.join(result)


def parse_and_execute(llm_output: str, hooks: HookRegistry = None) -> tuple[str, bool]:
    """
    解析 LLM 输出，提取工具调用并执行。

    v2: 统一 action/tool 格式，支持内置工具 + macOS 工具
    策略：尝试所有 ```json 块（从第一个开始），找到第一个可执行的
    """
    # 尝试从 ```json ... ``` 代码块提取
    # 使用贪婪匹配找到最外层的 ```json...``` 对
    # 注意：内容中可能含有 ``` 标记（如 markdown 代码块），需要特殊处理
    matches = _extract_json_blocks(llm_output)
    if not matches:
        # fallback: 找内联 JSON
        matches = re.findall(r'\{[^{}]*(?:"action"|"tool")\s*:.*?\}', llm_output, re.DOTALL)
    if not matches:
        return llm_output, False

    # 尝试每个 match（从第一个开始），找到第一个包含 action/tool 的可解析 JSON
    data = None
    for candidate in matches:
        candidate = candidate.strip()
        parsed = _try_parse_json(candidate)
        if parsed and ("action" in parsed or "tool" in parsed):
            data = parsed
            break

    if data is None:
        return f"(无法解析 JSON: {matches[0][:100]})\n{llm_output}", False

    # 统一格式: action/tool 都支持
    tool_name = data.get("action", data.get("tool", ""))

    # 完成信号
    if tool_name in ("done", "terminate"):
        answer = data.get("summary", data.get("input", {}).get("answer", "")) if isinstance(data.get("input"), dict) else data.get("summary", "任务完成")
        return str(answer) if answer else "任务完成", True

    # 路由到工具
    tool = TOOL_REGISTRY.get(tool_name)
    if not tool:
        return f"❌ 未知工具: {tool_name}. 可用: {', '.join(sorted(TOOL_REGISTRY.keys())[:15])}...", False

    # 准备输入
    tool_input = data.get("input", {})
    if not tool_input:
        # 兼容旧格式: {"tool": "shell", "command": "ls"} → input = {"command": "ls"}
        tool_input = {k: v for k, v in data.items() if k not in ("action", "tool")}

    if isinstance(tool_input, dict):
        tool_input_str = json.dumps(tool_input, ensure_ascii=False)
    else:
        tool_input_str = str(tool_input)

    # Hook: pre_tool
    if hooks:
        hooks.fire("pre_tool", term=tool, input=tool_input_str, ctx=None)

    # 执行
    import time
    t0 = time.time()
    try:
        if hasattr(tool, 'apply'):
            result = tool.apply(tool_input_str)
        else:
            result = tool(tool_input_str)
    except Exception as e:
        result = f"[ERROR] {e}"
    duration_ms = (time.time() - t0) * 1000

    # Hook: post_tool
    if hooks:
        output_wrapper = {"value": result}
        hooks.fire("post_tool", term=tool, input=tool_input_str,
                   output=output_wrapper, duration_ms=duration_ms, ctx=None)
        result = output_wrapper["value"]

    return str(result), False


# ════════════════════════════════════════════════════════════
# PersonalAssistant v2
# ════════════════════════════════════════════════════════════

class PersonalAssistant:
    """
    lambda 个人助理 v2 = Memory(ReAct(Brain)) + 21 内置工具 + Hook

    新增能力:
      - 21 个内置工具 (文件/代码/Shell/Git/Web/Notebook/Task)
      - ToolGateway 安全网关
      - Hook 全局审计
      - ContextManager 上下文压缩
      - TerminalUI 流式显示
    """

    def __init__(self, model: str = "sonnet", use_api: bool = False,
                 backend: str = "", verbose: bool = False):
        # Brain (LLM)
        if backend == BACKEND_OLLAMA:
            self.brain = OllamaLam(
                "lambda_v2", prompt=SYSTEM_PROMPT,
                model=model, temperature=0.3, max_tokens=4096,
            )
        elif use_api or backend == BACKEND_API:
            self.brain = Lam(
                "lambda_v2", prompt=SYSTEM_PROMPT,
                model=model, temperature=0.3, max_tokens=4096,
            )
        else:
            self.brain = ClaudeLam(
                "lambda_v2", prompt=SYSTEM_PROMPT,
                model=model, max_tokens=4096,
            )

        self.ctx = Context()
        self.conversation_history: list[dict] = []
        self.max_history = 30

        # v2 新增
        self.hooks = HookRegistry()
        self.context_manager = ContextManager(max_tokens=100000, keep_recent=5)
        self.ui = TerminalUI(verbose=verbose)

        # 注册审计 hook
        self._total_tokens = 0
        self._total_tool_calls = 0
        self.hooks.register("post_tool", self._audit_tool)

    def _audit_tool(self, **kwargs):
        """全局工具调用审计"""
        self._total_tool_calls += 1

    def _build_input(self, user_msg: str, observations: list[str] | None = None) -> str:
        """构建完整输入（历史 + 消息 + 工具结果）"""
        parts = []

        # 对话历史
        if self.conversation_history:
            parts.append("=== 对话历史 ===")
            for entry in self.conversation_history[-self.max_history:]:
                content = entry["content"]
                if len(content) > 300:
                    content = content[:300] + "..."
                parts.append(f"[{entry['role']}] {content}")
            parts.append("=== 历史结束 ===\n")

        # 工具观察
        if observations:
            parts.append("=== 工具执行结果 ===")
            for obs in observations:
                parts.append(obs)
            parts.append("=== 结果结束 ===\n")
            parts.append("请基于以上结果决定下一步。你的工具已验证可用。直接输出一个JSON代码块调用下一个工具。绝不要说'工具不可用'或让用户手动操作。")
        else:
            parts.append(f"[用户] {user_msg}")

        state = "\n".join(parts)

        # v2: 上下文压缩
        if self.context_manager.should_compact(state):
            state = self.context_manager.compact(state)

        return state

    def chat(self, user_msg: str) -> str:
        """
        与 lambda 对话。完整 ReAct 循环 (v2 流式交互)。
        """
        self.conversation_history.append({"role": "用户", "content": user_msg})

        observations = []
        max_steps = 15
        final_response = ""

        import time as _time

        for step in range(max_steps):
            # 显示当前步骤状态
            _step_label = f"[步骤 {step + 1}/{max_steps}]"

            # 构建输入
            if step == 0:
                llm_input = self._build_input(user_msg)
            else:
                llm_input = self._build_input(user_msg, observations)

            # 显示思考状态（如果不是流式模式，显示 spinner）
            is_streaming = hasattr(self.brain, 'stream') and self.brain.stream
            if not is_streaming:
                print(f"  ⏳ {_step_label} 思考中...", end="\r", flush=True)

            # Brain 思考 (β-规约) — 流式模式下 ClaudeLam 会实时输出
            t0 = _time.time()
            llm_output = self.brain.apply(llm_input, self.ctx)
            think_ms = (_time.time() - t0) * 1000

            if not is_streaming:
                # 清除 spinner
                print(f"  ✅ {_step_label} 思考完成 ({think_ms:.0f}ms)          ")

            # 解析并执行 (v2: 支持内置工具 + hooks)
            result, is_done = parse_and_execute(str(llm_output), self.hooks)

            if is_done:
                text_part = re.sub(r'```json.*?```', '', str(llm_output), flags=re.DOTALL).strip()
                final_response = text_part + ("\n" + result if result else "")
                print(f"  ✅ {_step_label} 任务完成")
                break

            if result != str(llm_output):
                # 工具被执行
                text_part = re.sub(r'```json.*?```', '', str(llm_output), flags=re.DOTALL).strip()
                json_matches = re.findall(r'\{[^{}]*(?:"action"|"tool")\s*:.*?\}', str(llm_output))
                tool_call = json_matches[-1] if json_matches else ""

                tool_name = ""
                try:
                    tool_name = json.loads(tool_call).get("action", json.loads(tool_call).get("tool", ""))
                except Exception:
                    pass

                # 显示思考内容
                if text_part and not is_streaming:
                    print(f"  💭 {text_part[:200]}")

                # 显示工具调用（带进度）
                t0_tool = _time.time()
                print(f"  🔧 {_step_label} 调用工具: {tool_name or 'unknown'}")

                # 工具输入预览
                try:
                    tool_input_data = json.loads(tool_call).get("input", {})
                    if isinstance(tool_input_data, dict):
                        input_preview = json.dumps(tool_input_data, ensure_ascii=False)[:120]
                    else:
                        input_preview = str(tool_input_data)[:120]
                    print(f"     输入: {input_preview}")
                except Exception:
                    pass

                tool_ms = (_time.time() - t0_tool) * 1000

                # 工具结果预览
                result_preview = result[:300].replace("\n", "\n     ")
                print(f"     结果: {result_preview}")
                print()

                observations.append(f"[工具调用] {tool_call}\n[执行结果] {result}")
            else:
                # 纯文本回复，没有工具调用
                # 检查是否应该继续（用户任务可能还没完成）
                output_str = str(llm_output)
                has_pending_action = any(kw in user_msg for kw in [
                    "保存", "写入", "创建文件", "save", "write",
                    "生成报告", "导出", "保存到", "写到",
                ])

                if has_pending_action and step < max_steps - 2 and not any(
                    sig in output_str.lower() for sig in ["任务完成", "已保存", "已写入", "已创建"]
                ):
                    # 用户要求了文件操作但 LLM 只输出了文本，继续推动
                    observations.append(
                        f"[LLM回复] {output_str[:500]}\n"
                        f"[系统提醒] 你已生成了内容但还没有保存到文件。"
                        f"请立即调用WriteFile工具将内容写入用户要求的路径。"
                    )
                    print(f"  💭 {_step_label} 已生成内容，继续执行文件保存...")
                else:
                    # 真的是最终回复
                    final_response = output_str
                    break
        else:
            final_response = f"(达到最大步数 {max_steps})\n最后的观察:\n" + (
                observations[-1] if observations else "无"
            )

        self.conversation_history.append({
            "role": "lambda",
            "content": final_response[:500],
        })

        return final_response

    def print_trace(self):
        """打印 β-规约追踪 (v2: 使用 TerminalUI)"""
        if self.ctx.trace:
            self.ui.trace(self.ctx.trace)
        else:
            print("  (无追踪记录)")

    def print_stats(self):
        """打印会话统计"""
        print(f"\n📊 会话统计:")
        print(f"  β-规约步数: {len(self.ctx.trace)}")
        print(f"  工具调用数: {self._total_tool_calls}")
        print(f"  对话轮数: {len(self.conversation_history) // 2}")