ai-robot-core/ai-service/app/models/mid/tool_trace.py

"""
Tool trace models for Mid Platform.
[AC-IDMP-15] 工具调用结构化记录

Reference: spec/intent-driven-mid-platform/openapi.provider.yaml - ToolCallTrace
"""

import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any


class ToolCallStatus(str, Enum):
    """工具调用状态"""
    OK = "ok"
    TIMEOUT = "timeout"
    ERROR = "error"
    REJECTED = "rejected"


class ToolType(str, Enum):
    """工具类型"""
    INTERNAL = "internal"
    MCP = "mcp"


@dataclass
class ToolCallTrace:
    """
    [AC-IDMP-15] 工具调用追踪记录
    Reference: openapi.provider.yaml - ToolCallTrace
    
    记录字段：
    - tool_name: 工具名称
    - tool_type: 工具类型 (internal | mcp)
    - registry_version: 注册表版本
    - auth_applied: 是否应用鉴权
    - duration_ms: 调用耗时（毫秒）
    - status: 调用状态 (ok | timeout | error | rejected)
    - error_code: 错误码
    - args_digest: 参数摘要（脱敏）
    - result_digest: 结果摘要
    - arguments: 完整参数
    - result: 完整结果
    """
    tool_name: str
    duration_ms: int
    status: ToolCallStatus
    tool_type: ToolType = ToolType.INTERNAL
    registry_version: str | None = None
    auth_applied: bool = False
    error_code: str | None = None
    args_digest: str | None = None
    result_digest: str | None = None
    arguments: dict[str, Any] | None = None
    result: Any = None
    started_at: datetime = field(default_factory=datetime.utcnow)
    completed_at: datetime | None = None

    def to_dict(self) -> dict[str, Any]:
        result = {
            "tool_name": self.tool_name,
            "duration_ms": self.duration_ms,
            "status": self.status.value,
        }
        if self.tool_type != ToolType.INTERNAL:
            result["tool_type"] = self.tool_type.value
        if self.registry_version:
            result["registry_version"] = self.registry_version
        if self.auth_applied:
            result["auth_applied"] = self.auth_applied
        if self.error_code:
            result["error_code"] = self.error_code
        if self.args_digest:
            result["args_digest"] = self.args_digest
        if self.result_digest:
            result["result_digest"] = self.result_digest
        if self.arguments:
            result["arguments"] = self.arguments
        if self.result is not None:
            result["result"] = self.result
        return result

    @staticmethod
    def compute_digest(data: Any, max_length: int = 64) -> str:
        """
        计算数据摘要（用于脱敏记录）
        
        Args:
            data: 原始数据
            max_length: 最大长度限制
            
        Returns:
            摘要字符串
        """
        if data is None:
            return ""
        
        if isinstance(data, (dict, list)):
            data_str = json.dumps(data, ensure_ascii=False, sort_keys=True)
        else:
            data_str = str(data)
        
        if len(data_str) <= max_length:
            return data_str
        
        hash_value = hashlib.sha256(data_str.encode("utf-8")).hexdigest()[:16]
        preview = data_str[:32]
        return f"{preview}...[hash:{hash_value}]"


@dataclass
class ToolCallBuilder:
    """
    [AC-IDMP-15] 工具调用记录构建器
    用于在工具执行过程中逐步构建追踪记录
    """
    tool_name: str
    tool_type: ToolType = ToolType.INTERNAL
    registry_version: str | None = None
    auth_applied: bool = False
    _started_at: datetime = field(default_factory=datetime.utcnow)
    _args: Any = None
    _result: Any = None
    _error: Exception | None = None
    _status: ToolCallStatus = ToolCallStatus.OK
    _error_code: str | None = None

    def with_args(self, args: Any) -> "ToolCallBuilder":
        """设置调用参数"""
        self._args = args
        return self

    def with_registry_info(self, version: str, auth_applied: bool) -> "ToolCallBuilder":
        """设置注册表信息"""
        self.registry_version = version
        self.auth_applied = auth_applied
        return self

    def with_result(self, result: Any) -> "ToolCallBuilder":
        """设置调用结果"""
        self._result = result
        self._status = ToolCallStatus.OK
        return self

    def with_error(self, error: Exception, error_code: str | None = None) -> "ToolCallBuilder":
        """设置错误信息"""
        self._error = error
        self._error_code = error_code
        if isinstance(error, TimeoutError):
            self._status = ToolCallStatus.TIMEOUT
        else:
            self._status = ToolCallStatus.ERROR
        return self

    def with_rejected(self, reason: str) -> "ToolCallBuilder":
        """设置拒绝状态"""
        self._status = ToolCallStatus.REJECTED
        self._error_code = reason
        return self

    def build(self) -> ToolCallTrace:
        """构建追踪记录"""
        completed_at = datetime.utcnow()
        duration_ms = int((completed_at - self._started_at).total_seconds() * 1000)
        
        return ToolCallTrace(
            tool_name=self.tool_name,
            tool_type=self.tool_type,
            registry_version=self.registry_version,
            auth_applied=self.auth_applied,
            duration_ms=duration_ms,
            status=self._status,
            error_code=self._error_code,
            args_digest=ToolCallTrace.compute_digest(self._args) if self._args else None,
            result_digest=ToolCallTrace.compute_digest(self._result) if self._result else None,
            started_at=self._started_at,
            completed_at=completed_at,
        )