agent_management/app.py

"""
FastAPI Web服务 - AI Agent管理服务
支持平台Agent和自定义Agent两种类型
"""
from fastapi import FastAPI, HTTPException, Depends
from pydantic import BaseModel, Field
from typing import Dict, List, Optional
from sqlalchemy.orm import Session
from datetime import datetime
import logging

from k8s_manager import K8sManager
from database import (
    get_db, Template, Agent, Quota, AgentMetric,
    AgentType, AgentStatus, parse_resource_string
)
import os

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(
    title="AI Agent Manager",
    description="Kubernetes AI Agent管理服务",
    version="1.0.0"
)

# 初始化K8s管理器
NAMESPACE = os.getenv("NAMESPACE", "ai-agents")
KUBECONFIG_PATH = os.getenv("KUBECONFIG_PATH", None)  # 可选：指定kubeconfig路径
k8s_manager = K8sManager(namespace=NAMESPACE, kubeconfig_path=KUBECONFIG_PATH)


# ==================== 请求/响应模型 ====================

# Template Management Models
class CreateTemplateRequest(BaseModel):
    """创建模板请求"""
    name: str = Field(..., min_length=1, max_length=100)
    display_name: str
    description: Optional[str] = None
    agent_type: str = Field(..., description="platform or custom")
    image: str
    port: Optional[int] = None
    env_requirements: Optional[Dict] = Field(default_factory=dict)
    cpu_request: Optional[str] = None
    cpu_limit: Optional[str] = None
    memory_request: Optional[str] = None
    memory_limit: Optional[str] = None
    min_replicas: int = 1
    max_replicas: int = 3
    target_cpu_utilization: int = 80


class UpdateTemplateRequest(BaseModel):
    """更新模板请求"""
    display_name: Optional[str] = None
    description: Optional[str] = None
    image: Optional[str] = None
    port: Optional[int] = None
    env_requirements: Optional[Dict] = None
    cpu_request: Optional[str] = None
    cpu_limit: Optional[str] = None
    memory_request: Optional[str] = None
    memory_limit: Optional[str] = None
    min_replicas: Optional[int] = None
    max_replicas: Optional[int] = None
    target_cpu_utilization: Optional[int] = None
    is_active: Optional[bool] = None


class TemplateResponse(BaseModel):
    """模板响应"""
    id: int
    name: str
    display_name: str
    description: Optional[str]
    agent_type: str
    image: str
    port: Optional[int]
    env_requirements: Dict
    cpu_request: Optional[str]
    cpu_limit: Optional[str]
    memory_request: Optional[str]
    memory_limit: Optional[str]
    min_replicas: int
    max_replicas: int
    target_cpu_utilization: int
    is_active: bool
    created_at: datetime
    
    class Config:
        from_attributes = True


# Platform Agent Models
class CreatePlatformAgentRequest(BaseModel):
    """创建平台Agent请求"""
    name: str = Field(..., min_length=1, max_length=63)
    template_name: str
    owner_id: str
    channel_id: Optional[str] = None
    tenant_id: Optional[str] = None
    query_params: Optional[Dict] = Field(default_factory=dict)


# Custom Agent Models
class ScalingConfig(BaseModel):
    """弹性伸缩配置"""
    min_replicas: int = Field(1, ge=0)
    max_replicas: int = Field(3, ge=1)
    target_cpu_utilization: int = Field(80, ge=1, le=100)


class CreateCustomAgentRequest(BaseModel):
    """创建自定义Agent请求"""
    name: str = Field(..., min_length=1, max_length=63)
    template_name: str
    owner_id: str
    channel_id: Optional[str] = None
    tenant_id: Optional[str] = None
    environment_vars: Dict[str, str]
    cpu_request: Optional[str] = None
    cpu_limit: Optional[str] = None
    memory_request: Optional[str] = None
    memory_limit: Optional[str] = None
    scaling_config: Optional[ScalingConfig] = None


class UpdateAgentEnvRequest(BaseModel):
    """更新Agent环境变量请求"""
    environment_vars: Dict[str, str]


class UpdateScalingRequest(BaseModel):
    """更新伸缩配置请求"""
    min_replicas: Optional[int] = None
    max_replicas: Optional[int] = None
    target_cpu_utilization: Optional[int] = None


# Unified Agent Response
class AgentResponseNew(BaseModel):
    """Agent响应（新）"""
    id: int
    name: str
    display_name: Optional[str]
    template_name: str
    agent_type: str
    status: str
    owner_id: str
    channel_id: Optional[str]
    tenant_id: Optional[str]
    service_url: Optional[str]
    current_replicas: int
    min_replicas: int
    max_replicas: int
    created_at: datetime
    last_accessed_at: Optional[datetime]
    
    class Config:
        from_attributes = True


# Legacy Models (for backward compatibility)
class CreateAgentRequest(BaseModel):
    """创建Agent请求（旧版）"""
    name: str = Field(..., description="Agent名称", min_length=1, max_length=63)
    template: str = Field(..., description="模板类型")
    config: Dict = Field(default_factory=dict, description="配置信息")
    env: Optional[Dict[str, str]] = Field(default_factory=dict, description="环境变量")


class AgentResponse(BaseModel):
    """Agent响应"""
    name: str
    namespace: str
    status: str
    created_at: Optional[str] = None
    template: Optional[str] = None
    service_port: Optional[int] = None
    access_info: Optional[Dict] = None
    pod_id: Optional[str] = None
    pod_ip: Optional[str] = None
    host_ip: Optional[str] = None
    node_name: Optional[str] = None
    owner_info: Optional[Dict] = None


class ResourceUsage(BaseModel):
    """资源使用情况"""
    cpu: Optional[str] = None
    memory: Optional[str] = None
    available: Optional[bool] = None
    reason: Optional[str] = None


class ResourceInfo(BaseModel):
    """资源信息（配额和使用情况）"""
    requests: Optional[Dict] = None
    limits: Optional[Dict] = None
    usage: Optional[ResourceUsage] = None


class PodStatusResponse(BaseModel):
    """Pod状态响应"""
    name: str
    namespace: str
    status: str
    template: Optional[str] = None
    created_at: Optional[str] = None
    node: Optional[str] = None
    pod_ip: Optional[str] = None
    resources: Optional[ResourceInfo] = None
    service_port: Optional[int] = None
    access_url: Optional[str] = None
    endpoints: Optional[Dict] = None
    conditions: Optional[List[Dict]] = None


class PodMetricsResponse(BaseModel):
    """Pod资源使用响应"""
    name: str
    requests: Dict
    limits: Dict


class MessageResponse(BaseModel):
    """通用消息响应"""
    status: str
    message: str


@app.get("/")
async def root():
    """健康检查"""
    return {
        "service": "AI Agent Manager",
        "status": "running",
        "namespace": NAMESPACE
    }


@app.post("/agents", response_model=AgentResponse)
async def create_agent(request: CreateAgentRequest):
    """
    创建AI Agent Pod
    
    Args:
        request: 创建请求（name, template, config, user_id可选）
    
    Returns:
        创建的Agent信息包括pod_id
    """
    try:
        logger.info(f"收到创建Agent请求: {request.name}, 模板: {request.template}")
        
        # 验证模板类型
        valid_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent"]
        if request.template not in valid_templates:
            raise HTTPException(
                status_code=400,
                detail=f"无效的模板类型。支持的模板: {', '.join(valid_templates)}"
            )
        
        # 合并环境变量到config
        config_data = request.config.copy()
        if request.env:
            config_data["env"] = request.env
            logger.info(f"环境变量: {list(request.env.keys())}")
        
        # 添加 user_id 标签
        user_id = config_data.get("user_id", "default")
        if "labels" not in config_data:
            config_data["labels"] = {}
        config_data["labels"]["user-id"] = user_id
        config_data["labels"]["managed-by"] = "agent-manager"
        
        # 创建Pod
        result = k8s_manager.create_pod(
            pod_name=request.name,
            template=request.template,
            config_data=config_data
        )
        
        # 获取 Pod 详细信息（包括 pod_id）
        try:
            import time
            time.sleep(1)  # 等待 Pod 创建完成
            pod = k8s_manager.v1.read_namespaced_pod(
                name=request.name,
                namespace=NAMESPACE
            )
            result["pod_id"] = pod.metadata.uid
            result["pod_ip"] = pod.status.pod_ip
            result["host_ip"] = pod.status.host_ip
            result["node_name"] = pod.spec.node_name
            result["owner_info"] = {
                "user_id": user_id,
                "agent_name": request.name,
                "namespace": NAMESPACE,
                "labels": pod.metadata.labels
            }
            logger.info(f"✅ Agent创建成功，Pod ID: {result['pod_id']}, 用户: {user_id}")
        except Exception as e:
            logger.warning(f"获取Pod详细信息失败: {str(e)}")
        
        return AgentResponse(**result)
    
    except Exception as e:
        logger.error(f"创建Agent失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/agents/{agent_name}", response_model=MessageResponse)
async def delete_agent(agent_name: str):
    """
    删除AI Agent Pod
    
    Args:
        agent_name: Agent名称
    
    Returns:
        删除结果
    """
    try:
        logger.info(f"收到删除Agent请求: {agent_name}")
        result = k8s_manager.delete_pod(pod_name=agent_name)
        
        if result.get("status") == "not_found":
            raise HTTPException(status_code=404, detail=result.get("message"))
        
        return MessageResponse(**result)
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"删除Agent失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/agents/{agent_name}/status", response_model=PodStatusResponse)
async def get_agent_status(agent_name: str):
    """
    获取Agent状态
    
    Args:
        agent_name: Agent名称
    
    Returns:
        Agent状态信息
    """
    try:
        logger.info(f"获取Agent状态: {agent_name}")
        result = k8s_manager.get_pod_status(pod_name=agent_name)
        
        if result.get("status") == "not_found":
            raise HTTPException(status_code=404, detail=result.get("message"))
        
        return PodStatusResponse(**result)
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取Agent状态失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/agents/{agent_name}/metrics", response_model=PodMetricsResponse)
async def get_agent_metrics(agent_name: str):
    """
    获取Agent资源使用情况
    
    Args:
        agent_name: Agent名称
    
    Returns:
        Agent资源使用信息
    """
    try:
        logger.info(f"获取Agent资源信息: {agent_name}")
        result = k8s_manager.get_pod_metrics(pod_name=agent_name)
        return PodMetricsResponse(**result)
    
    except Exception as e:
        logger.error(f"获取Agent资源信息失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/agents")
async def list_agents(template: Optional[str] = None):
    """
    列出所有Agent
    
    Args:
        template: 模板类型过滤（可选）
    
    Returns:
        Agent列表
    """
    try:
        logger.info(f"列出Agents, 模板过滤: {template}")
        
        label_selector = "managed-by=agent-manager"
        if template:
            label_selector += f",template={template}"
        
        result = k8s_manager.list_pods(label_selector=label_selector)
        return {"agents": result, "count": len(result)}
    
    except Exception as e:
        logger.error(f"列出Agents失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/templates")
async def list_templates():
    """
    列出所有可用的Agent模板及其所需参数
    
    Returns:
        模板列表及其配置信息
    """
    valid_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent"]
    
    templates_info = []
    for template in valid_templates:
        info = k8s_manager.get_template_info(template)
        templates_info.append(info)
    
    return {
        "templates": templates_info,
        "count": len(templates_info)
    }


@app.get("/templates/platform")
async def list_platform_templates():
    """
    获取平台 Agent 镜像列表
    
    Returns:
        平台提供的Agent模板列表
    """
    # 平台 Agent 是预定义的标准模板
    platform_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "jina_search_agent"]
    
    templates_info = []
    for template in platform_templates:
        info = k8s_manager.get_template_info(template)
        info["type"] = "platform"
        templates_info.append(info)
    
    return {
        "templates": templates_info,
        "count": len(templates_info),
        "type": "platform"
    }


@app.get("/templates/custom")
async def list_custom_templates():
    """
    获取自定义 Agent 镜像列表
    
    Returns:
        用户自定义的Agent模板列表
    """
    # 自定义 Agent 是用户可以配置数据库连接的模板
    custom_templates = ["mysql_agent", "postgresql_agent"]
    
    templates_info = []
    for template in custom_templates:
        info = k8s_manager.get_template_info(template)
        info["type"] = "custom"
        templates_info.append(info)
    
    return {
        "templates": templates_info,
        "count": len(templates_info),
        "type": "custom"
    }


@app.get("/templates/{template_name}")
async def get_template_info(template_name: str):
    """
    获取指定模板的详细信息
    
    Args:
        template_name: 模板名称
    
    Returns:
        模板详细信息（端口、所需环境变量等）
    """
    valid_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent"]
    
    if template_name not in valid_templates:
        raise HTTPException(
            status_code=404,
            detail=f"模板 {template_name} 不存在。可用模板: {', '.join(valid_templates)}"
        )
    
    return k8s_manager.get_template_info(template_name)


if __name__ == "__main__":
    import uvicorn
    
    host = os.getenv("SERVICE_HOST", "0.0.0.0")
    port = int(os.getenv("SERVICE_PORT", "8000"))
    
    logger.info(f"启动AI Agent Manager服务: {host}:{port}")
    uvicorn.run(app, host=host, port=port)