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

agent_runtime/runtime_engine/observability.py

@@ -0,0 +1,154 @@
+"""Langfuse/OpenTelemetry integration for the smolagents runtime."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+import logging
+import os
+from typing import Any, Iterator
+
+
+logger = logging.getLogger(__name__)
+
+_FALSE_VALUES = {"0", "false", "no", "off", "disable", "disabled"}
+_SMOLAGENTS_INSTRUMENTED = False
+
+
+@dataclass(frozen=True)
+class A2ATaskObservation:
+    agent_id: str
+    task_id: str
+    context_id: str | None
+    request_headers: dict[str, str]
+    task_input: str
+
+
+def setup_smolagents_observability() -> None:
+    """Enable the official smolagents OpenTelemetry instrumentation when configured."""
+
+    global _SMOLAGENTS_INSTRUMENTED
+    if _SMOLAGENTS_INSTRUMENTED or not _langfuse_enabled():
+        return
+
+    try:
+        from langfuse import get_client
+        from openinference.instrumentation.smolagents import SmolagentsInstrumentor
+    except ImportError:
+        logger.info(
+            "Langfuse env configured but smolagents telemetry dependencies are not "
+            "installed; runtime observability disabled"
+        )
+        return
+
+    try:
+        get_client()
+        SmolagentsInstrumentor().instrument()
+    except Exception:
+        logger.warning("Failed to initialize smolagents observability", exc_info=True)
+        return
+
+    _SMOLAGENTS_INSTRUMENTED = True
+    logger.info("Smolagents observability instrumentation enabled")
+
+
+@contextmanager
+def a2a_task_observation(
+    observation: A2ATaskObservation | None,
+) -> Iterator[Any | None]:
+    """Create a thin A2A task span and propagate task metadata to smolagents spans."""
+
+    if observation is None or not _langfuse_enabled():
+        yield None
+        return
+
+    try:
+        from langfuse import get_client, propagate_attributes
+    except ImportError:
+        yield None
+        return
+
+    langfuse = get_client()
+    metadata = _task_metadata(observation)
+    user_email = observation.request_headers.get("user_email") or None
+    session_id = observation.request_headers.get("X-Session-Id") or None
+
+    span_cm = langfuse.start_as_current_observation(
+        as_type="span",
+        name="a2a.task",
+        input={"text": observation.task_input},
+        metadata=metadata,
+    )
+    propagate_cm = propagate_attributes(
+        user_id=user_email,
+        session_id=session_id,
+        metadata=metadata,
+        trace_name="a2a.task",
+        tags=["a2a", observation.agent_id],
+    )
+
+    try:
+        with span_cm as span:
+            with propagate_cm:
+                try:
+                    yield span
+                except Exception as exc:
+                    _update_observation_error(span, exc)
+                    raise
+    finally:
+        _flush_langfuse(langfuse)
+
+
+def update_observation_output(span: Any | None, output: Any) -> None:
+    if span is None:
+        return
+    try:
+        span.update(output=output)
+    except Exception:
+        logger.debug("Failed to update Langfuse observation output", exc_info=True)
+
+
+def _langfuse_enabled() -> bool:
+    enabled = os.getenv("MCP_AGENT_LANGFUSE_ENABLED")
+    if enabled is not None and enabled.strip().lower() in _FALSE_VALUES:
+        return False
+    return all(
+        os.getenv(name)
+        for name in (
+            "LANGFUSE_PUBLIC_KEY",
+            "LANGFUSE_SECRET_KEY",
+            "LANGFUSE_BASE_URL",
+        )
+    )
+
+
+def _task_metadata(observation: A2ATaskObservation) -> dict[str, str]:
+    metadata = {
+        "agent_id": observation.agent_id,
+        "task_id": observation.task_id,
+        "context_id": observation.context_id or "",
+        "user_email": observation.request_headers.get("user_email", ""),
+        "session_id": observation.request_headers.get("X-Session-Id", ""),
+    }
+    return {key: value for key, value in metadata.items() if value}
+
+
+def _update_observation_error(span: Any, error: Exception) -> None:
+    try:
+        span.update(
+            output={
+                "error_type": type(error).__name__,
+                "message": str(error),
+            },
+            level="ERROR",
+            status_message=str(error),
+        )
+    except Exception:
+        logger.debug("Failed to update Langfuse observation error", exc_info=True)
+
+
+def _flush_langfuse(langfuse: Any) -> None:
+    try:
+        langfuse.flush()
+    except Exception:
+        logger.debug("Failed to flush Langfuse client", exc_info=True)

agent_runtime/runtime_engine/policy_registry.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+"""MCP 级工具治理的集中读写逻辑。
+
+这一层的目标是把 “tool policy 应该写到哪个 MCP 上” 这件事收敛到一个地方，
+避免 runtime 其他模块到处自己猜。
+"""
+
+from typing import Any
+
+from ..config.definitions import MCPSettings, ToolPolicy
+
+
+class MCPPolicyRegistry:
+    """围绕 MCPSettings.tool_policies 提供统一的查询与合并接口。"""
+
+    def __init__(
+        self, mcps: list[MCPSettings], tool_source_by_name: dict[str, str] | None = None
+    ):
+        self._mcps = mcps
+        self._tool_source_by_name = tool_source_by_name or {}
+
+    def find_mcp(self, mcp_name: str) -> MCPSettings | None:
+        return next((mcp for mcp in self._mcps if mcp.name == mcp_name), None)
+
+    def resolve_policy_mcp(
+        self, tool_name: str, mcp_name: str | None = None
+    ) -> MCPSettings | None:
+        if mcp_name is not None:
+            return self.find_mcp(mcp_name)
+
+        configured = [mcp for mcp in self._mcps if tool_name in mcp.tool_policies]
+        if len(configured) == 1:
+            return configured[0]
+        if len(configured) > 1:
+            return configured[0]
+
+        source_mcp = self._tool_source_by_name.get(tool_name)
+        if source_mcp:
+            return self.find_mcp(source_mcp)
+        if len(self._mcps) == 1:
+            return self._mcps[0]
+        return None
+
+    def set_tool_policy(
+        self, tool_name: str, policy: ToolPolicy, mcp_name: str | None = None
+    ) -> None:
+        target_mcp = self.resolve_policy_mcp(tool_name, mcp_name)
+        if target_mcp is None:
+            raise ValueError(f"cannot resolve MCP for tool '{tool_name}'")
+        target_mcp.tool_policies[tool_name] = policy
+
+    def get_tool_policy(
+        self, tool_name: str, mcp_name: str | None = None
+    ) -> ToolPolicy | None:
+        target_mcp = self.resolve_policy_mcp(tool_name, mcp_name)
+        if target_mcp is None:
+            return None
+        return target_mcp.tool_policies.get(tool_name)
+
+    def list_tool_policies(self, mcp_name: str | None = None) -> dict[str, ToolPolicy]:
+        if mcp_name is not None:
+            target_mcp = self.find_mcp(mcp_name)
+            return {} if target_mcp is None else dict(target_mcp.tool_policies)
+
+        policies: dict[str, ToolPolicy] = {}
+        for mcp in self._mcps:
+            for tool_name, policy in mcp.tool_policies.items():
+                policies[tool_name] = policy
+        return policies
+
+    def remove_tool_policy(self, tool_name: str, mcp_name: str | None = None) -> bool:
+        target_mcp = self.resolve_policy_mcp(tool_name, mcp_name)
+        if target_mcp is None:
+            return False
+        return target_mcp.tool_policies.pop(tool_name, None) is not None
+
+    def is_tool_enabled(self, tool_name: str, mcp_name: str | None = None) -> bool:
+        policy = self.get_tool_policy(tool_name, mcp_name)
+        return True if policy is None else policy.enabled
+
+    def merged_policy(
+        self, tool_name: str, decision: Any, mcp_name: str | None = None
+    ) -> ToolPolicy:
+        policy = (self.get_tool_policy(tool_name, mcp_name) or ToolPolicy()).model_copy(
+            deep=True
+        )
+        if decision is None:
+            return policy
+        if decision.requires_confirmation is not None:
+            policy.requires_confirmation = decision.requires_confirmation
+        if decision.input_fields:
+            policy.input_fields = decision.input_fields
+        if decision.prompt is not None:
+            policy.prompt = decision.prompt
+        if decision.allow_arg_override is not None:
+            policy.allow_arg_override = decision.allow_arg_override
+        return policy

agent_runtime/runtime_engine/protocol_tools.py

@@ -0,0 +1,94 @@
+"""运行时协议工具。
+
+这三个工具不是业务 MCP 工具，而是 runtime 和上层协议之间的桥梁：
+
+- `ask_user`: 主动向用户追问缺失信息
+- `ask_auth`: 向用户请求授权确认
+- `final_answer`: 统一输出最终结果
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .context import current_task_id, current_task_pool
+
+try:
+    from smolagents import Tool
+except (
+    ImportError
+):  # pragma: no cover - fallback for local static tests without smolagents
+
+    class Tool:  # type: ignore[no-redef]
+        name = ""
+        description = ""
+        inputs: dict[str, Any] = {}
+        output_type = "string"
+
+        def __init__(self):
+            pass
+
+
+def _consume_task_input(input_key: str) -> Any:
+    try:
+        task_id = current_task_id.get()
+        pool = current_task_pool.get()
+    except LookupError:
+        return ""
+
+    task_info = pool.get(task_id)
+    if not task_info or not hasattr(task_info, input_key):
+        return ""
+
+    value = getattr(task_info, input_key)
+    setattr(task_info, input_key, None)
+    return value or ""
+
+
+class AskUserTool(Tool):
+    name = "ask_user"
+    description = "向用户提出问题，补充信息。"
+    inputs = {
+        "task": {
+            "type": "string",
+            "description": "向用户表达你需要什么信息",
+        }
+    }
+    output_type = "string"
+
+    def forward(self, task: str):
+        return _consume_task_input("user_input")
+
+
+class AskAuthTool(Tool):
+    name = "ask_auth"
+    description = (
+        "向用户请求授权确认。用户端返回布尔值：true 表示确认，false 表示拒绝。"
+    )
+    inputs = {
+        "task": {
+            "type": "string",
+            "description": "向用户说明需要什么授权信息",
+        }
+    }
+    output_type = "string"
+
+    def forward(self, task: str):
+        return _consume_task_input("auth_required")
+
+
+class FinalAnswerTool(Tool):
+    name = "final_answer"
+    description = (
+        "agent 最后通过此工具输出最终结果，请确保最后一次调用的工具是 final_answer"
+    )
+    inputs = {
+        "answer": {
+            "type": "string",
+            "description": "最终结果内容",
+        }
+    }
+    output_type = "string"
+
+    def forward(self, answer: str):
+        return answer