taiji-AI-PAD/services/mcp-server/schemas.py

"""
Pydantic schemas for API requests and responses
"""

from datetime import datetime
from typing import Any, Dict, List, Optional, Union
from pydantic import BaseModel, Field, validator
import uuid


class BaseSchema(BaseModel):
    """基础schema类"""
    class Config:
        from_attributes = True
        json_encoders = {
            datetime: lambda v: v.isoformat(),
            uuid.UUID: lambda v: str(v),
        }


# ========== MCP协议相关 ==========

class MCPRequest(BaseModel):
    """MCP请求模型"""
    jsonrpc: str = "2.0"
    id: Union[str, int] = Field(default_factory=lambda: str(uuid.uuid4()))
    method: str
    params: Optional[Dict[str, Any]] = None


class MCPResponse(BaseModel):
    """MCP响应模型"""
    jsonrpc: str = "2.0"
    id: Union[str, int]
    result: Optional[Any] = None
    error: Optional[Dict[str, Any]] = None


class MCPError(BaseModel):
    """MCP错误模型"""
    code: int
    message: str
    data: Optional[Any] = None


# ========== 工具相关 ==========

class ToolParameter(BaseModel):
    """工具参数定义"""
    name: str
    type: str
    description: Optional[str] = None
    required: bool = True
    default: Optional[Any] = None
    enum: Optional[List[Any]] = None


class ToolDefinition(BaseSchema):
    """工具定义"""
    name: str
    description: str
    category: Optional[str] = None
    parameters: List[ToolParameter] = []
    returns: Optional[Dict[str, Any]] = None
    
    # API相关
    endpoint: Optional[str] = None
    method: str = "POST"
    headers: Optional[Dict[str, str]] = None
    
    # 限制信息
    rate_limit: int = 100
    timeout: int = 30
    cost_per_call: float = 0.0


class ToolExecution(BaseModel):
    """工具执行请求"""
    tool_name: str
    parameters: Dict[str, Any]
    timeout: Optional[int] = None


class ToolResult(BaseSchema):
    """工具执行结果"""
    success: bool
    result: Optional[Any] = None
    error: Optional[str] = None
    execution_time: float = 0.0
    cost: float = 0.0


# ========== Agent相关 ==========

class AgentCreateRequest(BaseModel):
    """创建Agent请求"""
    name: str = Field(..., min_length=1, max_length=100)
    description: Optional[str] = None
    role: str = Field(
        default="general-purpose agent",
        min_length=1,
        max_length=200
    )
    goal: str = Field(
        default="Handle generic MCP tasks and routing",
        min_length=1
    )
    
    tools: List[str] = []  # 工具名称列表
    config: Dict[str, Any] = {}
    capabilities: List[str] = []
    
    owner_id: Optional[uuid.UUID] = None
    
    @validator('name')
    def validate_name(cls, v):
        """验证Agent名称"""
        if not v.replace('-', '').replace('_', '').isalnum():
            raise ValueError('名称只能包含字母、数字、连字符和下划线')
        return v


class AgentUpdateRequest(BaseModel):
    """更新Agent请求"""
    name: Optional[str] = None
    description: Optional[str] = None
    role: Optional[str] = None
    goal: Optional[str] = None
    tools: Optional[List[str]] = None
    config: Optional[Dict[str, Any]] = None
    capabilities: Optional[List[str]] = None


class AgentCard(BaseSchema):
    """Agent卡片信息"""
    id: uuid.UUID
    name: str
    description: Optional[str]
    role: str
    goal: str
    
    tools: List[str] = []
    capabilities: List[str] = []
    
    # 端点信息
    endpoints: Dict[str, str] = {}
    
    # 状态信息
    status: str = "active"
    version: str = "1.0.0"
    
    # 统计信息
    total_executions: int = 0
    success_rate: float = 0.0
    avg_execution_time: float = 0.0
    
    # 时间信息
    created_at: datetime
    updated_at: datetime


class AgentExecution(BaseModel):
    """Agent执行请求"""
    method: str
    params: Optional[Dict[str, Any]] = None
    timeout: Optional[int] = None
    session_id: Optional[str] = None


class ExecutionResult(BaseSchema):
    """执行结果"""
    execution_id: str
    success: bool
    result: Optional[Any] = None
    error: Optional[str] = None
    
    # 性能指标
    execution_time: float = 0.0
    cpu_usage: float = 0.0
    memory_usage: float = 0.0
    network_io: float = 0.0
    
    # 成本信息
    eu_consumed: float = 0.0
    cost: float = 0.0
    
    # 时间戳
    started_at: datetime
    completed_at: Optional[datetime] = None


# ========== 用户相关 ==========

class UserCreate(BaseModel):
    """创建用户请求"""
    username: str = Field(..., min_length=3, max_length=50)
    email: str = Field(..., pattern=r'^[^@]+@[^@]+\.[^@]+$')
    password: str = Field(..., min_length=8)
    full_name: Optional[str] = None


class UserUpdate(BaseModel):
    """更新用户请求"""
    email: Optional[str] = None
    full_name: Optional[str] = None
    is_active: Optional[bool] = None


class UserResponse(BaseSchema):
    """用户响应"""
    id: uuid.UUID
    username: str
    email: str
    full_name: Optional[str]
    is_active: bool
    is_admin: bool
    created_at: datetime
    updated_at: datetime


class UserLogin(BaseModel):
    """用户登录请求"""
    username: str
    password: str


class Token(BaseModel):
    """访问令牌"""
    access_token: str
    token_type: str = "bearer"
    expires_in: int


# ========== 会话相关 ==========

class SessionCreate(BaseModel):
    """创建会话请求"""
    context: Optional[Dict[str, Any]] = None
    metadata: Optional[Dict[str, Any]] = None


class SessionResponse(BaseSchema):
    """会话响应"""
    id: uuid.UUID
    session_id: str
    status: str
    context: Dict[str, Any]
    metadata: Dict[str, Any]
    created_at: datetime
    updated_at: datetime


# ========== 计费相关 ==========

class BillingRecord(BaseSchema):
    """计费记录"""
    id: uuid.UUID
    execution_id: uuid.UUID
    eu_consumed: float
    cost: float
    currency: str
    
    # 资源详情
    cpu_time: float
    memory_max: float
    network_io: float
    storage_io: float
    
    created_at: datetime


class BillingSummary(BaseModel):
    """计费汇总"""
    user_id: uuid.UUID
    period_start: datetime
    period_end: datetime
    
    total_executions: int
    total_eu_consumed: float
    total_cost: float
    currency: str
    
    # 按服务分类
    breakdown_by_agent: Dict[str, float] = {}
    breakdown_by_tool: Dict[str, float] = {}


# ========== API响应包装 ==========

class APIResponse(BaseModel):
    """API响应包装"""
    success: bool
    message: str = ""
    data: Optional[Any] = None
    timestamp: datetime = Field(default_factory=datetime.utcnow)


class PaginatedResponse(BaseModel):
    """分页响应"""
    items: List[Any]
    total: int
    page: int
    page_size: int
    has_next: bool
    has_prev: bool


# ========== 系统状态 ==========

class HealthCheck(BaseModel):
    """健康检查响应"""
    status: str
    timestamp: datetime
    services: Dict[str, str]
    version: str = "1.0.0"


class SystemMetrics(BaseModel):
    """系统指标"""
    timestamp: datetime
    
    # 服务指标
    active_agents: int
    total_executions: int
    success_rate: float
    avg_response_time: float
    
    # 资源指标
    cpu_usage: float
    memory_usage: float
    disk_usage: float
    
    # 业务指标
    daily_active_users: int
    total_eu_consumed: float
    total_cost: float


# ========== 错误响应 ==========

class ErrorResponse(BaseModel):
    """错误响应"""
    error: str
    message: str
    details: Optional[Dict[str, Any]] = None
    timestamp: datetime = Field(default_factory=datetime.utcnow)


class ValidationError(BaseModel):
    """验证错误"""
    field: str
    message: str
    invalid_value: Optional[Any] = None