Tool_Agent/src/tool_agent/models.py

"""公共数据模型"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field


# ---- 枚举 ----

class ToolStatus(str, Enum):
    ACTIVE = "active"
    INACTIVE = "inactive"
    STAGING = "staging"
    BUILDING = "building"


class BackendRuntime(str, Enum):
    """工具后端执行环境（调用层统一为本地 Python/FastAPI）"""
    LOCAL = "local"    # 本地 Python 运行时（纯 uv 子进程）
    DOCKER = "docker"  # Docker 容器运行时
    REMOTE = "remote"  # 云端 API / 远程服务后端


class MessageType(str, Enum):
    TOOL_REQUEST = "tool_request"
    TOOL_READY = "tool_ready"
    TOOL_ERROR = "tool_error"
    HEALTH_ALERT = "health_alert"


class ContainerStatus(str, Enum):
    RUNNING = "running"
    STOPPED = "stopped"
    DESTROYED = "destroyed"


# ---- 工具元信息（纯元数据，不含运行时） ----

class ToolMeta(BaseModel):
    """工具元信息 — 只描述工具是什么，不管在哪跑"""
    model_config = {"extra": "ignore"}   # 兼容旧 registry.json 中的多余字段

    tool_id: str
    name: str
    category: str = ""  # 功能领域分类（保留向后兼容）
    description: str = ""
    input_schema: dict[str, Any] = Field(default_factory=dict)
    output_schema: dict[str, Any] = Field(default_factory=dict)
    stream_support: bool = False
    status: ToolStatus = ToolStatus.ACTIVE

    # 新增字段（向后兼容）
    backend_runtime: BackendRuntime | None = None  # 后端执行环境
    provider_ids: list[str] = Field(default_factory=list)  # 工具提供方
    capability_ids: list[str] = Field(default_factory=list)  # 工具解决的原子能力
    knowledge_ids: list[str] = Field(default_factory=list)  # 工具挂载的知识资产

# ---- 容器信息（Docker 运行时，独立于工具元数据） ----

class ContainerInfo(BaseModel):
    container_id: str
    tool_id: str = ""
    image: str = ""
    port_mapping: dict[int, int] = Field(default_factory=dict)
    volumes: dict[str, str] = Field(default_factory=dict)
    mem_limit: str = "1g"
    nano_cpus: int = 1_000_000_000
    use_gpu: bool = False
    gpu_count: int = -1
    status: ContainerStatus = ContainerStatus.RUNNING
    created_at: datetime | None = None
    last_accessed: datetime | None = None
    destroyed_at: datetime | None = None


# ---- 请求/响应 ----

class InvokeRequest(BaseModel):
    params: dict[str, Any] = Field(default_factory=dict)
    stream: bool = False


class InvokeResponse(BaseModel):
    status: str = "success"
    result: Any = None
    error: str | None = None


class ToolRequest(BaseModel):
    """外部提交新工具需求"""
    description: str
    callback_url: str | None = None


class TaskStatus(BaseModel):
    task_id: str
    status: str = "pending"
    result: Any = None


# ---- 内部消息 ----

class AgentMessage(BaseModel):
    type: MessageType
    payload: dict[str, Any] = Field(default_factory=dict)