agent-runtime-sdk 0.1.0__py3-none-any.whl

agent_runtime/runtime_engine/tool_flow.py

@@ -0,0 +1,332 @@
+"""工具调用前的运行时治理逻辑。
+
+注意：这里不真正执行工具，只在“工具即将执行”前做几件事：
+
+- 归一化参数
+- 检查是否缺少必填输入
+- 必要时向用户追问
+- 必要时向用户请求授权确认
+- 调用 plugin 的 before_tool / after_tool 钩子
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any, Callable
+
+from .context import (
+    AuthRequiredWaitState,
+    InputRequiredWaitState,
+    TaskContext,
+    TaskPhase,
+    WaitState,
+    current_task_id,
+)
+from ..config.definitions import ToolPolicy
+from ..exceptions import TaskCancelledError, TaskWaitTimeoutError, UserCancelledError
+from ..plugins import ToolExecutionContext
+from ..config.policies import (
+    collect_missing_fields,
+    merge_input_fields,
+    parse_user_payload,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeToolFlow:
+    """负责工具执行前后的治理钩子与等待中断逻辑。"""
+
+    def _build_tool_callbacks(
+        self, tools: list[Any], extra_tools: list[Any]
+    ) -> dict[str, list[Callable]]:
+        tool_callbacks: dict[str, list[Callable]] = {
+            "ask_user": [self._interrupt_before_ask_user],
+            "ask_auth": [self._interrupt_before_ask_auth],
+        }
+        for tool in [*tools, *extra_tools]:
+            tool_name = getattr(tool, "name", "")
+            if tool_name and tool_name not in {"ask_user", "ask_auth", "final_answer"}:
+                tool_callbacks[tool_name] = [self._build_tool_callback(tool_name)]
+        return tool_callbacks
+
+    def _extract_tool_prompt(self, chat_message: Any, tool_name: str) -> str | None:
+        for tool_call in chat_message.tool_calls or []:
+            if tool_call.function.name != tool_name:
+                continue
+            args = tool_call.function.arguments
+            if isinstance(args, dict):
+                return args.get("task") or args.get("prompt") or args.get("question")
+            return str(args) if args is not None else None
+        return None
+
+    def _wait_for_event(self, task_id: str) -> None:
+        task_info = self.task_pool.get(task_id)
+        if not task_info:
+            return
+        event = task_info.event
+        loop = task_info.loop
+        if event and loop:
+            future = asyncio.run_coroutine_threadsafe(event.wait(), loop)
+            future.result()
+
+    def _current_task(self) -> tuple[str | None, TaskContext | None]:
+        try:
+            task_id = current_task_id.get()
+        except LookupError:
+            return None, None
+        return task_id, self.task_pool.get(task_id)
+
+    @staticmethod
+    def _coerce_tool_arguments(raw_args: Any) -> dict[str, Any]:
+        if isinstance(raw_args, dict):
+            return dict(raw_args)
+        if isinstance(raw_args, str):
+            try:
+                parsed = json.loads(raw_args)
+            except Exception:
+                return {"value": raw_args}
+            return parsed if isinstance(parsed, dict) else {"value": parsed}
+        return {}
+
+    @staticmethod
+    def _set_wait_item(task_info: TaskContext, wait_item: WaitState) -> None:
+        if isinstance(wait_item, InputRequiredWaitState):
+            phase = TaskPhase.WAITING_INPUT
+        else:
+            phase = TaskPhase.WAITING_AUTH
+        task_info.set_wait(wait_item, phase)
+
+    @staticmethod
+    def _raise_if_wait_timed_out(task_info: TaskContext) -> None:
+        if task_info.timed_out:
+            raise TaskWaitTimeoutError(task_info.timeout_reason or "任务等待超时")
+
+    @staticmethod
+    def _raise_if_task_stopped(task_info: TaskContext) -> None:
+        if task_info.control.stop_requested:
+            raise task_info.control.stop_error or TaskCancelledError(
+                "task cancelled while waiting"
+            )
+
+    def _collect_required_input(
+        self,
+        *,
+        task_id: str,
+        task_info: TaskContext,
+        tool_name: str,
+        args: dict[str, Any],
+        policy: ToolPolicy,
+    ) -> dict[str, Any]:
+        missing_fields = collect_missing_fields(policy, args)
+        while missing_fields:
+            prompt = policy.prompt or f"请补充工具 {tool_name} 所需信息"
+            logger.debug(
+                "Task waiting for user input task_id=%s tool=%s missing_fields=%s",
+                task_id,
+                tool_name,
+                [field.name for field in missing_fields],
+            )
+            self._set_wait_item(task_info, InputRequiredWaitState(prompt=prompt))
+            self._wait_for_event(task_id)
+            self._raise_if_task_stopped(task_info)
+            self._raise_if_wait_timed_out(task_info)
+            resume_payload = parse_user_payload(task_info.user_input)
+            task_info.user_input = None
+            args = merge_input_fields(args, policy, resume_payload)
+            missing_fields = collect_missing_fields(policy, args)
+        return args
+
+    def _request_tool_confirmation(
+        self,
+        *,
+        task_id: str,
+        task_info: TaskContext,
+        tool_name: str,
+        args: dict[str, Any],
+        policy: ToolPolicy,
+    ) -> Any:
+        if not policy.requires_confirmation:
+            return args
+
+        if task_info.wait_item is not None:
+            return args
+
+        logger.debug(
+            "Task waiting for auth confirmation task_id=%s tool=%s",
+            task_id,
+            tool_name,
+        )
+        self._set_wait_item(
+            task_info, AuthRequiredWaitState(tool_name=tool_name, args=args)
+        )
+        self._wait_for_event(task_id)
+        self._raise_if_task_stopped(task_info)
+        self._raise_if_wait_timed_out(task_info)
+        if task_info.auth_denied:
+            task_info.auth_denied = True
+            raise UserCancelledError("user denied auth request")
+
+        if task_info.tool_args_override is not None:
+            override = task_info.tool_args_override
+            override_name = task_info.tool_name_override
+            task_info.tool_args_override = None
+            task_info.tool_name_override = None
+            if policy.allow_arg_override and (
+                not override_name or override_name == tool_name
+            ):
+                return override
+        return args
+
+    def _prepare_tool_arguments(
+        self,
+        *,
+        tool_name: str,
+        mcp_name: str | None,
+        args: Any,
+        task_id: str,
+        task_info: TaskContext,
+        tool_call: Any,
+    ) -> Any:
+        plugin = self.plugin
+        if plugin is None:
+            return args
+
+        context = ToolExecutionContext(
+            agent_definition=self.definition,
+            tool_name=tool_name,
+            mcp_name=mcp_name,
+            args=dict(args),
+            task_info=task_info,
+            raw_tool_call=tool_call,
+        )
+
+        self._raise_if_task_stopped(task_info)
+
+        normalized = plugin.normalize_args(context)
+        if isinstance(normalized, dict):
+            args = normalized
+            context.args = dict(args)
+
+        decision = plugin.before_tool(context)
+        policy = self._merged_policy(tool_name, decision, mcp_name)
+        args = self._collect_required_input(
+            task_id=task_id,
+            task_info=task_info,
+            tool_name=tool_name,
+            args=args,
+            policy=policy,
+        )
+        args = self._request_tool_confirmation(
+            task_id=task_id,
+            task_info=task_info,
+            tool_name=tool_name,
+            args=args,
+            policy=policy,
+        )
+        context.args = (
+            dict(args) if isinstance(args, dict) else self._coerce_tool_arguments(args)
+        )
+        return args
+
+    def _build_tool_callback(self, tool_name: str):
+        source_mcp = self._tool_source_by_name.get(tool_name)
+
+        def _callback(chat_message, agent=None):
+            tool_call = self._find_tool_call(chat_message, tool_name)
+            if tool_call is None:
+                return
+
+            task_id, task_info = self._current_task()
+            if task_id is None or task_info is None:
+                return
+
+            logger.debug(
+                "Preparing tool call task_id=%s tool=%s mcp=%s",
+                task_id,
+                tool_name,
+                source_mcp,
+            )
+            args = self._coerce_tool_arguments(tool_call.function.arguments)
+            tool_call.function.arguments = self._prepare_tool_arguments(
+                tool_name=tool_name,
+                mcp_name=source_mcp,
+                args=args,
+                task_id=task_id,
+                task_info=task_info,
+                tool_call=tool_call,
+            )
+
+        return _callback
+
+    def _merged_policy(
+        self, tool_name: str, decision: Any, mcp_name: str | None = None
+    ) -> ToolPolicy:
+        return self._policy_registry.merged_policy(tool_name, decision, mcp_name)
+
+    @staticmethod
+    def _find_tool_call(chat_message: Any, tool_name: str):
+        for tool_call in chat_message.tool_calls or []:
+            if tool_call.function.name == tool_name:
+                return tool_call
+        return None
+
+    def _interrupt_before_ask_user(self, chat_message, agent=None):
+        task_id, task_info = self._current_task()
+        if task_id is None or task_info is None:
+            return
+        logger.debug("Task entered explicit ask_user wait task_id=%s", task_id)
+        self._set_wait_item(
+            task_info,
+            InputRequiredWaitState(
+                prompt=self._extract_tool_prompt(chat_message, "ask_user")
+                or "需要您补充信息"
+            ),
+        )
+        self._wait_for_event(task_id)
+        self._raise_if_task_stopped(task_info)
+        self._raise_if_wait_timed_out(task_info)
+
+    def _interrupt_before_ask_auth(self, chat_message, agent=None):
+        task_id, task_info = self._current_task()
+        if task_id is None or task_info is None:
+            return
+        logger.debug("Task entered explicit ask_auth wait task_id=%s", task_id)
+        self._set_wait_item(task_info, AuthRequiredWaitState())
+        self._wait_for_event(task_id)
+        self._raise_if_task_stopped(task_info)
+        self._raise_if_wait_timed_out(task_info)
+        if task_info.auth_denied:
+            task_info.auth_denied = True
+            raise UserCancelledError("user denied auth request")
+
+    def _after_tool_hook(self, tool_output):
+        if self.plugin is None:
+            return
+        tool_name = (
+            getattr(tool_output.tool_call, "name", "") if tool_output.tool_call else ""
+        )
+        source_mcp = self._tool_source_by_name.get(tool_name)
+        task_id, task_info = self._current_task()
+        if task_id is not None:
+            logger.debug(
+                "Finished tool call task_id=%s tool=%s mcp=%s",
+                task_id,
+                tool_name,
+                source_mcp,
+            )
+        if task_info is not None:
+            task_info.clear_active_tool()
+        self.plugin.after_tool(
+            ToolExecutionContext(
+                agent_definition=self.definition,
+                tool_name=tool_name,
+                mcp_name=source_mcp,
+                args={},
+                observation=tool_output.observation,
+                result=tool_output.output,
+            )
+        )