521 lines
15 KiB
Python
521 lines
15 KiB
Python
"""
|
||
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
|
||
namespace: Optional[str] = None
|
||
requests: Dict
|
||
limits: Dict
|
||
usage: Optional[Dict] = None # 实时使用情况（需要 metrics-server）
|
||
timestamp: Optional[str] = None # metrics 时间戳
|
||
metrics_available: Optional[bool] = None # metrics-server 是否可用
|
||
|
||
|
||
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)
|