precis-cli 0.1.3__py3-none-any.whl

app/shared/services/ai/agent/executor.py

@@ -0,0 +1,554 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Copyright 2026 Precis Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""@fileoverview Agent 执行器
+
+功能概述:
+- 实现 Agent 主循环：LLM 调用 → 工具执行 → 结果回传 → 下一轮
+- 支持最大迭代轮数、取消检查、checkpoint 回调
+- 约定：名为 generate_config 的工具是最终输出工具，调用成功并返回 config 时任务结束
+
+架构设计:
+- AgentExecutor 组合 AIProvider、ToolRegistry、AgentMemory
+- 每轮调用 LLM，若返回 tool_calls 则执行工具
+- 若工具调用中包含 generate_config 且成功返回 config，则以该 config 作为最终结果结束
+- 若 LLM 直接返回可解析为配置 JSON 的文本，也接受为最终结果
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from app.shared.services.llm.providers.base import BaseProvider, ChatMessage, ChatRequest
+
+from .memory import AgentMemory
+from .tool_registry import ToolRegistry
+from .types import AgentMetrics, AgentResult, AgentTurn, ToolResult
+
+logger = logging.getLogger(__name__)
+
+
+def _tool_call_signature(tc: Any) -> str:
+    """提取 tool_call 的去重签名。
+
+    用于防止同一动作被重复执行（Ollama 多行重复场景）。
+    - id 非空：用 id（OpenAI 等提供稳定 id 的 provider 走快路径）
+    - id 为空：用 function.name + arguments 的组合签名（Ollama 的 id 通常为空）
+    """
+    if not isinstance(tc, dict):
+        return ""
+    tc_id: str = str(tc.get("id", ""))
+    if tc_id:
+        return tc_id
+    func = tc.get("function", {})
+    if not isinstance(func, dict):
+        return ""
+    name: str = str(func.get("name", ""))
+    args: str = str(func.get("arguments", ""))
+    return f"{name}:{args}"
+
+
+def _cross_turn_signature(tc: Any) -> str:
+    """跨轮重复检测专用签名：只看 name + arguments，不看 id。
+
+    与 _tool_call_signature 的区别：跨轮场景下 id 必然不同（LLM 每次新生成），
+    若用 id 会导致"同工具同参数"的重复调用无法被识别。本函数专注于"做了同样的事"。
+
+    用于 P1-2 收敛引导的跨轮重复检测。
+    """
+    if not isinstance(tc, dict):
+        return ""
+    func = tc.get("function", {})
+    if not isinstance(func, dict):
+        return ""
+    name: str = str(func.get("name", ""))
+    args: str = str(func.get("arguments", ""))
+    return f"{name}:{args}"
+
+
+# 默认系统提示词
+DEFAULT_SYSTEM_PROMPT = """你是一个数据治理专家 Agent，擅长分析数据文件并生成数据验证配置。
+你可以调用工具来完成任务。每一步请根据观察结果决定下一步行动。
+当需要输出最终结果时，请调用 generate_config 工具，其返回的 config 会被作为最终结果。"""
+
+# 最终输出工具名
+FINAL_OUTPUT_TOOL = "generate_config"
+
+
+# 防僵死总时长下限（秒）：即使单次 IO 超时很小，整体守卫也不低于 5 分钟
+STREAM_GUARD_MIN_TIMEOUT = 300.0
+# 单次 IO 超时 → 整体防僵死上限的放大倍数：多轮工具调用的长回复需要远大于单次 IO 的预算
+STREAM_GUARD_MULTIPLIER = 5
+
+
+def stream_guard_timeout(provider_timeout: Any) -> float:
+    """由 provider 单次 IO 超时推导 Agent 流式调用的整体防僵死上限。
+
+    两层超时的分工：
+    - provider 内部的 aiohttp connect/read 超时（timeout_seconds，Ollama 默认 60s）
+      管"单次 IO"——一次读不到数据就快速失败，避免单个请求挂死；
+    - 本层 wait_for 只防"整体僵死"（流永不结束、取消信号失效），不限制生成时长。
+
+    若直接复用 timeout_seconds 作总上限，本地模型多轮长回复会被 60s 掐断；
+    故取 timeout_seconds * 5 且不低于 300s。非实数（mock/未实现 provider）回退 300s。
+
+    返回:
+        整体防僵死上限（秒）
+    """
+    # 必须校验为实数，否则非数值（如 MagicMock）参与运算会产生无意义/非法的 timeout
+    if isinstance(provider_timeout, (int, float)):
+        return max(provider_timeout * STREAM_GUARD_MULTIPLIER, STREAM_GUARD_MIN_TIMEOUT)
+    return STREAM_GUARD_MIN_TIMEOUT
+
+
+# 进度回调签名：stage, progress(0-1), extra dict
+def default_progress_callback(stage: str, progress: float, extra: dict[str, Any] | None = None) -> None:
+    pass
+
+
+# 默认取消回调
+def default_cancelled_callback() -> bool:
+    return False
+
+
+# 默认 tool result 回调
+def default_on_tool_result(tr: ToolResult) -> None:
+    pass
+
+
+class AgentExecutor:
+    """
+    @classdesc Agent 执行器
+
+    实现规划-执行-观察的循环。
+    """
+
+    def __init__(
+        self,
+        provider: BaseProvider,
+        registry: ToolRegistry,
+        system_prompt: str | None = None,
+        max_iterations: int = 5,
+        max_tokens: int = 120000,
+        progress_callback: Callable[[str, float, dict[str, Any] | None], None] | None = None,
+        checkpoint_callback: Callable[[dict[str, Any]], None] | None = None,
+        cancelled_callback: Callable[[], bool] | None = None,
+        on_tool_result: Callable[[ToolResult], None] | None = None,
+        on_chunk: Callable[[str], None] | None = None,
+        on_turn: Callable[[int], None] | None = None,
+        on_tool_call: Callable[[str, str, int], None] | None = None,
+        final_output_tool: str | None = FINAL_OUTPUT_TOOL,
+    ):
+        """
+        @methoddesc 初始化 Agent 执行器
+
+        参数:
+            provider: AIProvider 配置
+            registry: 工具注册表
+            system_prompt: 系统提示词，None 使用默认
+            max_iterations: 最大迭代轮数
+            max_tokens: 记忆最大 token 预算
+            progress_callback: 进度回调
+            checkpoint_callback: checkpoint 保存回调
+            cancelled_callback: 取消检查回调
+            on_tool_result: 工具结果回调，用于提取 metrics
+            on_chunk: 流式文本增量回调（逐字触发），用于 SSE 流式输出
+            on_turn: 新轮次开始回调，参数为轮次序号
+            on_tool_call: 工具调用开始回调，参数为 (tool_name, call_id, turn)
+        """
+        self.provider = provider
+        self.registry = registry
+        self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
+        self.max_iterations = max_iterations
+        self.max_tokens = max_tokens
+        self.progress_callback = progress_callback or default_progress_callback
+        self.checkpoint_callback = checkpoint_callback or (lambda x: None)
+        self.cancelled_callback = cancelled_callback or default_cancelled_callback
+        self.on_tool_result = on_tool_result or default_on_tool_result
+        self.on_chunk = on_chunk or (lambda text: None)
+        self.on_turn = on_turn or (lambda turn: None)
+        self.on_tool_call = on_tool_call or (lambda name, call_id, turn: None)
+        # 最终输出工具名：generation agent 用 "generate_config"，chat agent 传 None 跳过最终配置提取
+        self.final_output_tool = final_output_tool
+        self._chat_service: Any | None = None
+        self._memory: AgentMemory | None = None
+        # P1-2 跨轮 tool_call 重复检测：记录上一轮的工具调用签名，本轮命中时注入提醒
+        self._prev_tool_signatures: set[str] = set()
+
+    async def run(
+        self,
+        task_message: str,
+        initial_checkpoint: dict[str, Any] | None = None,
+    ) -> AgentResult:
+        """
+        @methoddesc 运行 Agent
+
+        参数:
+            task_message: 用户任务描述
+            initial_checkpoint: 可选的初始 checkpoint 数据,含 messages 时从断点续跑
+
+        返回:
+            AgentResult 执行结果
+        """
+        # 续跑:从 checkpoint 恢复 memory,而非新建
+        if initial_checkpoint and initial_checkpoint.get("messages"):
+            self._memory = AgentMemory.restore_from_checkpoint(initial_checkpoint)
+            start_turn = initial_checkpoint.get("turn", 0) + 1
+        else:
+            self._memory = AgentMemory(self.system_prompt, max_tokens=self.max_tokens)
+            self._memory.set_task(task_message)
+            start_turn = 1
+
+        # P1-2：每次 run 重置跨轮签名集合，避免跨任务污染
+        self._prev_tool_signatures = set()
+
+        tools = self.registry.get_definitions()
+
+        result = AgentResult(success=True, iterations=0)
+
+        for turn_idx in range(start_turn, self.max_iterations + 1):
+            # 取消检查点 1: turn 开始
+            if self.cancelled_callback():
+                result.success = False
+                result.cancelled = True
+                result.error = "任务已取消"
+                return result
+
+            self.on_turn(turn_idx)
+            self.progress_callback(
+                f"agent_turn_{turn_idx}",
+                turn_idx / self.max_iterations,
+                {"iterations": turn_idx, "max_iterations": self.max_iterations},
+            )
+
+            # P1-2 收敛引导：在最后一轮（达到 max_iterations）提醒 LLM 直接回复，
+            # 避免它继续调工具导致循环用尽后直接 fail。reminder 在 get_messages 前注入。
+            if turn_idx >= self.max_iterations:
+                self._memory.add_system_reminder(
+                    f"注意：你已用尽全部 {self.max_iterations} 轮工具调用预算。"
+                    "请立即基于已获取的信息给出最终回复，不要再调用任何工具。"
+                )
+
+            messages = self._memory.get_messages()
+            logger.debug(f"Agent turn {turn_idx}, messages={len(messages)}")
+
+            chat_messages = []
+            for m in messages:
+                chat_messages.append(
+                    ChatMessage(
+                        role=m["role"],
+                        content=m.get("content"),
+                        tool_calls=m.get("tool_calls"),
+                        tool_call_id=m.get("tool_call_id"),
+                    )
+                )
+            req = ChatRequest(
+                messages=chat_messages,
+                model=self.provider.model if hasattr(self.provider, "model") else None,
+                temperature=0.3,
+                tools=tools,
+                tool_choice="auto",
+            )
+
+            # 流式调用: 累积文本 + tool_calls，逐字触发 on_chunk
+            # chat_stream 是 async generator function（async def + yield），直接 async for 即可
+            content = ""
+            raw_tool_calls: list[dict[str, Any]] = []
+            cancelled_in_stream = False
+            received_any_chunk = False  # 追踪是否收到过任何 chunk（区分"空流"与"返回但内容空"）
+            try:
+                # 用 wait_for 包裹整体流式读取，防止 provider 阻塞时取消信号永不触发。
+                # 两层超时分工：provider 内部的 connect/read 超时管单次 IO；
+                # 本层只防整体僵死，取单次超时的 5 倍且不低于 300s（见 stream_guard_timeout），
+                # 不限制生成长度——本地模型多轮长回复不会被 60s 单次超时连坐掐断。
+                stream_timeout = stream_guard_timeout(getattr(self.provider, "timeout_seconds", None))
+
+                async def _consume_stream() -> None:
+                    nonlocal content, raw_tool_calls, cancelled_in_stream, received_any_chunk
+                    async for chunk in self.provider.chat_stream(req):
+                        received_any_chunk = True
+                        # 取消检查点 2: 流内 chunk 间
+                        if self.cancelled_callback():
+                            cancelled_in_stream = True
+                            return
+                        if chunk.type == "delta":
+                            content += chunk.text or ""
+                            self.on_chunk(chunk.text or "")
+                        elif chunk.type == "tool_calls" and chunk.tool_calls:
+                            # 去重：防止 Ollama 多行重复导致同一动作多次执行。
+                            # 关键：Ollama 的 tool_call id 通常为空，旧守卫 `if tc_id and ...`
+                            # 对空 id 短路失效，恰在它要防的 Ollama 场景上漏过。
+                            # 修复：id 非空用 id（快路径）；否则用 name+arguments 签名兜底。
+                            for tc in chunk.tool_calls:
+                                sig = _tool_call_signature(tc)
+                                if sig and any(_tool_call_signature(existing) == sig for existing in raw_tool_calls):
+                                    continue
+                                raw_tool_calls.append(tc)
+
+                await asyncio.wait_for(_consume_stream(), timeout=stream_timeout)
+            except TimeoutError:
+                logger.error(f"Agent LLM 流式调用超时({stream_timeout}s)")
+                result.success = False
+                result.error = f"AI 服务响应超时({stream_timeout}s)"
+                return result
+            except Exception as e:
+                logger.error(f"Agent LLM 流式调用失败: {e}")
+                result.success = False
+                result.error = f"AI 服务调用失败: {e}"
+                return result
+
+            # 流内取消：保留已累积的 content 到 result（让 cancelled 事件携带 partial reply）
+            if cancelled_in_stream:
+                result.success = False
+                result.cancelled = True
+                result.error = "任务已取消"
+                result.content = content  # 保留部分内容
+                return result
+
+            # 空流检测：provider 没有返回任何 chunk（连空的 delta/tool_calls 都没有）
+            # 这才是真正的异常空流，避免注入非法空 assistant 消息导致后续轮次 provider 400 错误
+            # 注意：content 为空但有 tool_calls（或收到过空 delta）是合法的，不算空流
+            if not received_any_chunk:
+                logger.warning(f"Agent 第 {turn_idx} 轮 LLM 返回空响应（无任何 chunk）")
+                result.success = False
+                result.error = "LLM 返回空响应，请检查模型配置或重试"
+                return result
+
+            tool_calls = [self.registry.parse_tool_call(tc) for tc in raw_tool_calls]
+
+            # P1-2 跨轮重复检测：若本轮存在与上一轮完全相同的 tool_call（同工具同参数），
+            # 为下一轮注入提醒，避免 LLM 反复调同一工具同参数空转。
+            # 注意：reminder 追加到 memory 末尾，下一轮 get_messages 时 LLM 可见。
+            self._check_repeated_tool_calls(raw_tool_calls)
+
+            self._memory.add_assistant_message(content=content or None, tool_calls=raw_tool_calls or None)
+
+            turn = AgentTurn(
+                turn=turn_idx,
+                content=content,
+                tool_calls=tool_calls,
+            )
+
+            if tool_calls:
+                # 触发工具调用开始回调
+                for tc in tool_calls:
+                    self.on_tool_call(tc.name, tc.id, turn_idx)
+                # 执行工具
+                self.progress_callback("agent_executing_tools", turn_idx / self.max_iterations, None)
+                tool_results = await self.registry.execute_many(tool_calls)
+                turn.tool_results = tool_results
+
+                final_config = self._extract_final_config(tool_results)
+                if final_config is not None:
+                    # 最终输出工具已生成配置，任务结束
+                    result.config = final_config
+                    result.success = True
+                    result.iterations = turn_idx
+                    result.content = content
+                    self._memory.add_turn(turn)
+                    result.turns.append(turn)
+                    final_checkpoint = self._make_checkpoint(turn_idx, turn, result.metrics.to_dict())
+                    try:
+                        self.checkpoint_callback(final_checkpoint)
+                    except Exception:
+                        logger.warning("checkpoint_callback 异常，已忽略", exc_info=True)
+                    return result
+
+                for tr in tool_results:
+                    self.on_tool_result(tr)
+                    observation = tr.observation
+                    if isinstance(observation, dict):
+                        observation_text = self._dict_to_text(observation)
+                    else:
+                        observation_text = str(observation)
+                    # 失败时把 error 回灌给 LLM，否则 tool 消息 content 会是空串，
+                    # LLM 无法得知失败原因（apply_actions 的 did-you-mean 建议等也会丢失）。
+                    if not tr.success and tr.error:
+                        observation_text = (observation_text + "\n" if observation_text else "") + f"[ERROR] {tr.error}"
+                    self._memory.add_tool_result(tr.call_id, tr.name, observation_text)
+
+                self._memory.add_turn(turn)
+                result.turns.append(turn)
+                result.iterations = turn_idx
+
+                # 更新 checkpoint
+                checkpoint = self._make_checkpoint(turn_idx, turn, None)
+                try:
+                    self.checkpoint_callback(checkpoint)
+                except Exception:
+                    logger.warning("checkpoint_callback 异常，已忽略", exc_info=True)
+                continue
+
+            # 没有 tool_calls 且 content 非空，视为最终结果
+            if content:
+                self._memory.add_turn(turn)
+                result.turns.append(turn)
+                result.iterations = turn_idx
+                result.content = content
+
+                # 尝试从 content 解析配置
+                config = self._try_parse_config(content)
+                if config:
+                    result.config = config
+
+                final_checkpoint = self._make_checkpoint(turn_idx, turn, result.metrics.to_dict())
+                try:
+                    self.checkpoint_callback(final_checkpoint)
+                except Exception:
+                    logger.warning("checkpoint_callback 异常，已忽略", exc_info=True)
+                break
+
+            # content 为空且无 tool_calls，记录 turn 后继续
+            self._memory.add_turn(turn)
+            result.turns.append(turn)
+            result.iterations = turn_idx
+
+        # 如果达到最大轮数仍未结束。
+        # 注意：最后一轮以纯文本收敛（上方 break 的正常路径）时 content 非空且 config 为 None，
+        # 这是 chat agent 的合法结局，不能误判为失败——只有"既无 config 也无文本产出"才算用尽预算。
+        if result.iterations >= self.max_iterations and result.config is None and not result.content:
+            result.success = False
+            result.error = f"达到最大迭代轮数 {self.max_iterations}，仍未获得最终结果"
+
+        return result
+
+    def _check_repeated_tool_calls(self, raw_tool_calls: list[dict[str, Any]]) -> None:
+        """P1-2 跨轮重复调用检测。
+
+        比对本轮 tool_calls 与上一轮的签名集合：若存在完全相同的调用（同工具同参数），
+        为下一轮注入提醒，避免 LLM 反复调用同一工具同一参数空转浪费预算。
+        无论是否命中，都把本轮签名更新到集合，供下一轮比对。
+
+        参数:
+            raw_tool_calls: 本轮 LLM 返回的原始 tool_calls（OpenAI dict 格式）
+        """
+        # _memory 类型为 AgentMemory | None，但在 run() 中已保证调用本方法前完成初始化；
+        # 此处防御性判断同时为 mypy 提供 None 收敛，与 _create_checkpoint 中的写法保持一致。
+        if not self._memory or not raw_tool_calls:
+            return
+        current_signatures = {_cross_turn_signature(tc) for tc in raw_tool_calls}
+        # 去掉空签名（无法识别的 tool_call 不参与检测）
+        current_signatures.discard("")
+        if not current_signatures:
+            return
+
+        repeated = current_signatures & self._prev_tool_signatures
+        if repeated:
+            # 构造友好提示，列出重复的工具名（从签名解析，格式 "name:args"）
+            names = sorted({sig.split(":", 1)[0] for sig in repeated if ":" in sig})
+            tool_list = "、".join(names) if names else "同一工具"
+            self._memory.add_system_reminder(
+                f"提示：你刚才已经调用过 {tool_list}（相同参数），请勿重复调用同一工具同一参数。"
+                "请直接使用已有的工具结果继续推理，或在必要时更换参数/工具。"
+            )
+
+        # 更新签名集合供下一轮比对
+        self._prev_tool_signatures = current_signatures
+
+    def _extract_final_config(self, tool_results: list[ToolResult]) -> dict[str, Any] | None:
+        """从工具结果中提取最终配置（final_output_tool 指定工具的返回值）。
+
+        若 final_output_tool 为 None（如 chat agent），直接返回 None，跳过最终配置提取。
+        """
+        if not self.final_output_tool:
+            return None
+        for tr in tool_results:
+            if tr.name != self.final_output_tool:
+                continue
+            if not tr.success:
+                continue
+            observation = tr.observation
+            if isinstance(observation, dict) and observation.get("success"):
+                config = observation.get("config")
+                if isinstance(config, dict):
+                    return config
+        return None
+
+    def _make_checkpoint(
+        self,
+        turn: int,
+        agent_turn: AgentTurn,
+        metrics: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        """构造 checkpoint（含完整 messages，支持续跑）。"""
+        memory_checkpoint = self._memory.create_checkpoint() if self._memory else {}
+        return {
+            **memory_checkpoint,
+            "turn": turn,
+            "content": agent_turn.content,
+            "tool_calls": [{"id": c.id, "name": c.name, "arguments": c.arguments} for c in agent_turn.tool_calls],
+            "tool_results": [
+                {
+                    "call_id": r.call_id,
+                    "name": r.name,
+                    "success": r.success,
+                    "observation": r.observation if isinstance(r.observation, dict) else str(r.observation),
+                    "error": r.error,
+                }
+                for r in agent_turn.tool_results
+            ],
+            "metrics": metrics,
+        }
+
+    @staticmethod
+    def _dict_to_text(data: dict[str, Any]) -> str:
+        """将字典转为文本。"""
+        import json
+
+        try:
+            return json.dumps(data, ensure_ascii=False)
+        except (TypeError, ValueError):
+            return str(data)
+
+    @staticmethod
+    def _try_parse_config(content: str | None) -> dict[str, Any] | None:
+        """尝试从文本中提取 JSON 配置。"""
+        if not content:
+            return None
+        import json
+        import re
+
+        # 尝试找到 Markdown JSON 代码块
+        match = re.search(r"```(?:json)?\s*(.*?)\s*```", content, re.DOTALL)
+        text = match.group(1).strip() if match else content
+
+        start = text.find("{")
+        if start == -1:
+            return None
+        balance = 0
+        for i, char in enumerate(text[start:], start=start):
+            if char == "{":
+                balance += 1
+            elif char == "}":
+                balance -= 1
+                if balance == 0:
+                    try:
+                        parsed: Any = json.loads(text[start : i + 1])
+                        if isinstance(parsed, dict):
+                            return parsed
+                    except json.JSONDecodeError:
+                        return None
+        return None
+
+    @staticmethod
+    def build_metrics_from_issues(issues: list[dict[str, Any]]) -> AgentMetrics:
+        """根据校验问题构造 metrics。"""
+        metrics = AgentMetrics()
+        metrics.issues = issues
+        metrics.failed_rules = len(issues)
+        return metrics