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

agent_runtime/runtime_engine/agent_build.py

@@ -0,0 +1,315 @@
+"""单次请求的 agent 构建引擎。
+
+当一条新请求到来时，这里负责把下面几样东西组装成一个可执行 agent：
+
+- model
+- 当前请求可用的 MCP tools
+- runtime 内置协议工具
+- plugin 提供的额外工具
+- framework/plugin/instructions
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from typing import Any
+import weakref
+
+from ..mcp.discovery import load_mcp_client_and_tools
+from ..mcp.metadata import DiscoveredTool
+from ..mcp.utils import merge_mcp_headers, register_tool_source
+from ..exceptions import AgentBuildError, AgentRuntimeError, MCPToolLoadError
+from ..plugins import ToolExecutionContext
+from ..plugins.instructions import (
+    build_framework_instructions,
+    join_instruction_sections,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeAgentBuild:
+    """负责把一次请求需要的资源组装成可执行 agent。"""
+
+    @staticmethod
+    def _make_disconnect_cancel_hook(client: Any):
+        """构造一个不会强持有 client 的取消钩子。
+
+        这里不能直接把 `client.disconnect` 挂到 tool 上。
+        否则会形成：
+
+        - client -> tools
+        - tool -> bound method
+        - bound method -> client
+
+        的运行时对象环。
+        """
+
+        disconnect = getattr(type(client), "disconnect", None)
+        if disconnect is None:
+            return None
+        try:
+            client_ref = weakref.ref(client)
+        except TypeError:
+            return None
+
+        def _cancel() -> None:
+            resolved_client = client_ref()
+            if resolved_client is None:
+                return
+            disconnect(resolved_client)
+
+        return _cancel
+
+    def build_agent(self, mcp_headers: dict[str, str] | None = None):
+        """按单次请求构建 agent。"""
+        if self.plugin is None:
+            self.reload()
+        if self.plugin is None:
+            logger.error(
+                "Plugin unavailable while building agent agent_id=%s load_error=%s",
+                self._agent.agent_id,
+                self.load_error,
+            )
+            raise AgentBuildError(
+                self.load_error
+                or f"agent {self._agent.agent_id} plugin is not available",
+                user_message="Agent 构建失败，请检查启动日志",
+            )
+
+        model = self._build_model()
+        mcp_clients, tools = self._load_request_tools(mcp_headers or {})
+        try:
+            extra_tools = self._collect_extra_tools()
+            instructions = self._build_agent_instructions()
+            all_tools = self._compose_agent_tools(tools, extra_tools)
+            agent = self._create_agent_instance(
+                model=model,
+                all_tools=all_tools,
+                tools=tools,
+                extra_tools=extra_tools,
+                instructions=instructions,
+            )
+            self._attach_request_resources(agent, mcp_clients)
+            logger.debug(
+                "Built agent instance agent_id=%s mcp_tools=%s extra_tools=%s",
+                self._agent.agent_id,
+                len(tools),
+                len(extra_tools),
+            )
+            return agent
+        except Exception as exc:
+            logger.exception("Failed to build agent agent_id=%s", self._agent.agent_id)
+            for client in mcp_clients:
+                try:
+                    client.disconnect()
+                except Exception:
+                    logger.warning(
+                        "Failed to disconnect MCP client while rolling back agent build agent_id=%s",
+                        self._agent.agent_id,
+                        exc_info=True,
+                    )
+            if isinstance(exc, AgentRuntimeError):
+                raise
+            raise AgentBuildError(
+                f"failed to build agent '{self._agent.agent_id}': {exc}",
+            ) from exc
+
+    def _build_model(self, model_cls=None):
+        model_settings = self._runtime_config.model
+        if model_cls is not None:
+            return model_cls(
+                api_base=model_settings.api_base or "",
+                model_id=model_settings.model_id or "",
+                api_key=os.getenv(model_settings.api_key_env, ""),
+                flatten_messages_as_text=model_settings.flatten_messages_as_text,
+            )
+        from ..model_registry import ModelRegistry
+
+        return ModelRegistry.create(
+            model_settings.provider,
+            api_base=model_settings.api_base or "",
+            model_id=model_settings.model_id or "",
+            api_key=os.getenv(model_settings.api_key_env, ""),
+            flatten_messages_as_text=model_settings.flatten_messages_as_text,
+        )
+
+    def _load_request_tools(
+        self, request_headers: dict[str, str]
+    ) -> tuple[list[Any], list[Any]]:
+        clients: list[Any] = []
+        enabled_tools: list[Any] = []
+        loaded_tool_sources: dict[str, str] = {}
+
+        try:
+            for mcp in self._mcps:
+                if not mcp.enabled:
+                    logger.debug(
+                        "Skipping disabled MCP during runtime load agent_id=%s mcp=%s",
+                        self._agent.agent_id,
+                        mcp.name,
+                    )
+                    continue
+                logger.debug(
+                    "Loading MCP tools agent_id=%s mcp=%s url=%s",
+                    self._agent.agent_id,
+                    mcp.name,
+                    mcp.url,
+                )
+                client, tools = load_mcp_client_and_tools(
+                    mcp,
+                    headers=self._request_mcp_headers(mcp, request_headers),
+                )
+                clients.append(client)
+
+                for tool in tools:
+                    tool_name = getattr(tool, "name", "")
+                    if not tool_name or not self._is_tool_enabled(tool_name):
+                        continue
+                    register_tool_source(loaded_tool_sources, tool_name, mcp.name)
+                    try:
+                        tool._runtime_mcp_name = mcp.name
+                    except Exception:
+                        pass
+                    try:
+                        tool._runtime_cancel_hook = self._make_disconnect_cancel_hook(
+                            client
+                        )
+                    except Exception:
+                        pass
+                    enabled_tools.append(tool)
+            logger.debug(
+                "Loaded enabled MCP tools agent_id=%s tool_count=%s",
+                self._agent.agent_id,
+                len(enabled_tools),
+            )
+            return clients, enabled_tools
+        except Exception as exc:
+            logger.exception(
+                "Failed to load MCP tools agent_id=%s", self._agent.agent_id
+            )
+            for client in clients:
+                try:
+                    client.disconnect()
+                except Exception:
+                    logger.warning(
+                        "Failed to disconnect MCP client during tool load rollback agent_id=%s",
+                        self._agent.agent_id,
+                        exc_info=True,
+                    )
+            if isinstance(exc, MCPToolLoadError):
+                raise
+            raise MCPToolLoadError(
+                f"failed to load MCP tools for agent '{self._agent.agent_id}': {exc}",
+            ) from exc
+
+    def _collect_extra_tools(self) -> list[Any]:
+        if self.plugin is None:
+            return []
+        return self.plugin.extra_tools(self.definition)
+
+    def _build_agent_instructions(self) -> str:
+        plugin_instructions = ""
+        if self.plugin is not None:
+            plugin_instructions = self.plugin.build_instructions(
+                self.definition, self.discovered_tools
+            )
+        return join_instruction_sections(
+            build_framework_instructions(self.definition, self.discovered_tools),
+            self._agent.extra_instructions,
+            plugin_instructions,
+        )
+
+    @staticmethod
+    def _compose_agent_tools(tools: list[Any], extra_tools: list[Any]) -> list[Any]:
+        from .protocol_tools import AskAuthTool, AskUserTool, FinalAnswerTool
+
+        return [*tools, AskUserTool(), AskAuthTool(), FinalAnswerTool(), *extra_tools]
+
+    def _create_agent_instance(
+        self,
+        *,
+        model: Any,
+        all_tools: list[Any],
+        tools: list[Any],
+        extra_tools: list[Any],
+        instructions: str,
+    ):
+        from ..smol import ToolCallingCheckAgent
+
+        return ToolCallingCheckAgent(
+            model=model,
+            tools=all_tools,
+            prompt_templates=None,
+            instructions=instructions,
+            tool_callbacks=self._build_tool_callbacks(tools, extra_tools),
+            max_steps=self._runtime_config.model.max_steps,
+            stream_outputs=False,
+            name=_safe_agent_name(self._agent.agent_id),
+            description=self._agent.description,
+            after_tool_hook=self._after_tool_hook,
+        )
+
+    @staticmethod
+    def _attach_request_resources(agent: Any, mcp_clients: list[Any]) -> None:
+        agent._runtime_mcp_clients = mcp_clients
+
+    def format_result(self, result: Any) -> Any:
+        if self.plugin is None:
+            return result
+        formatted = self.plugin.format_result(
+            ToolExecutionContext(
+                agent_definition=self.definition,
+                tool_name="final_answer",
+                args={},
+                result=result,
+            )
+        )
+        return result if formatted is None else formatted
+
+    def list_tools(self, mcp_name: str | None = None) -> list[DiscoveredTool]:
+        """返回 discovery 阶段的工具摘要，不返回可执行 tool 实例。"""
+
+        if mcp_name is None:
+            return list(self.discovered_tools)
+        return [tool for tool in self.discovered_tools if tool.source_mcp == mcp_name]
+
+    def get_tool(
+        self, tool_name: str, mcp_name: str | None = None
+    ) -> DiscoveredTool | None:
+        """按名字获取 discovery 阶段工具摘要。"""
+
+        if mcp_name is None:
+            return next((t for t in self.discovered_tools if t.name == tool_name), None)
+        return next(
+            (
+                t
+                for t in self.discovered_tools
+                if t.name == tool_name and t.source_mcp == mcp_name
+            ),
+            None,
+        )
+
+    def _request_mcp_headers(
+        self, mcp, request_headers: dict[str, str]
+    ) -> dict[str, str]:
+        passed_headers = {
+            header_name: request_headers.get(header_name, "")
+            for header_name in mcp.pass_through_headers
+        }
+        return merge_mcp_headers(mcp, passed_headers)
+
+    def _is_tool_enabled(self, tool_name: str) -> bool:
+        return self._policy_registry.is_tool_enabled(tool_name)
+
+
+def _safe_agent_name(value: str) -> str:
+    normalized = re.sub(r"[^0-9a-zA-Z_]", "_", value)
+    if not normalized:
+        return "agent"
+    if normalized[0].isdigit():
+        normalized = f"agent_{normalized}"
+    return normalized