199 lines
5.3 KiB
Python
199 lines
5.3 KiB
Python
"""
|
|
Agent 服务
|
|
处理 Agent 相关的业务逻辑
|
|
"""
|
|
from typing import Optional, List
|
|
from datetime import datetime
|
|
from app.models import db, Agent, Gateway
|
|
|
|
|
|
class AgentService:
|
|
"""Agent 服务"""
|
|
|
|
def register_agent(
|
|
self,
|
|
name: str,
|
|
gateway_id: str = None,
|
|
model: str = None,
|
|
capabilities: list = None,
|
|
priority: int = 5,
|
|
weight: int = 10,
|
|
connection_limit: int = 5
|
|
) -> Agent:
|
|
"""注册新 Agent"""
|
|
agent = Agent(
|
|
name=name,
|
|
display_name=name,
|
|
gateway_id=gateway_id,
|
|
model=model,
|
|
capabilities=capabilities or [],
|
|
priority=priority,
|
|
weight=weight,
|
|
connection_limit=connection_limit,
|
|
status='offline'
|
|
)
|
|
|
|
db.session.add(agent)
|
|
|
|
# 更新 Gateway 的 Agent 计数
|
|
if gateway_id:
|
|
gateway = Gateway.query.get(gateway_id)
|
|
if gateway:
|
|
gateway.agent_count += 1
|
|
|
|
db.session.commit()
|
|
|
|
return agent
|
|
|
|
def get_agent(self, agent_id: str) -> Optional[Agent]:
|
|
"""获取 Agent"""
|
|
return Agent.query.get(agent_id)
|
|
|
|
def get_agents(
|
|
self,
|
|
status: str = None,
|
|
gateway_id: str = None
|
|
) -> List[Agent]:
|
|
"""获取 Agent 列表"""
|
|
query = Agent.query
|
|
|
|
if status:
|
|
query = query.filter_by(status=status)
|
|
|
|
if gateway_id:
|
|
query = query.filter_by(gateway_id=gateway_id)
|
|
|
|
return query.all()
|
|
|
|
def get_available_agents(self) -> List[Agent]:
|
|
"""获取可用 Agent 列表"""
|
|
return Agent.query.filter_by(status='online').filter(
|
|
Agent.current_sessions < Agent.connection_limit
|
|
).all()
|
|
|
|
def update_agent_status(
|
|
self,
|
|
agent_id: str,
|
|
status: str,
|
|
socket_id: str = None
|
|
) -> Optional[Agent]:
|
|
"""更新 Agent 状态"""
|
|
agent = Agent.query.get(agent_id)
|
|
if not agent:
|
|
return None
|
|
|
|
agent.status = status
|
|
|
|
if socket_id:
|
|
agent.socket_id = socket_id
|
|
|
|
if status == 'online':
|
|
agent.last_heartbeat = datetime.utcnow()
|
|
|
|
db.session.commit()
|
|
|
|
return agent
|
|
|
|
def heartbeat(self, agent_id: str) -> Optional[Agent]:
|
|
"""Agent 心跳"""
|
|
agent = Agent.query.get(agent_id)
|
|
if not agent:
|
|
return None
|
|
|
|
agent.last_heartbeat = datetime.utcnow()
|
|
agent.status = 'online'
|
|
|
|
db.session.commit()
|
|
|
|
return agent
|
|
|
|
def update_agent_config(
|
|
self,
|
|
agent_id: str,
|
|
name: str = None,
|
|
display_name: str = None,
|
|
model: str = None,
|
|
capabilities: list = None,
|
|
priority: int = None,
|
|
weight: int = None,
|
|
connection_limit: int = None
|
|
) -> Optional[Agent]:
|
|
"""更新 Agent 配置"""
|
|
agent = Agent.query.get(agent_id)
|
|
if not agent:
|
|
return None
|
|
|
|
if name is not None:
|
|
agent.name = name
|
|
if display_name is not None:
|
|
agent.display_name = display_name
|
|
if model is not None:
|
|
agent.model = model
|
|
if capabilities is not None:
|
|
agent.capabilities = capabilities
|
|
if priority is not None:
|
|
agent.priority = priority
|
|
if weight is not None:
|
|
agent.weight = weight
|
|
if connection_limit is not None:
|
|
agent.connection_limit = connection_limit
|
|
|
|
db.session.commit()
|
|
|
|
return agent
|
|
|
|
def delete_agent(self, agent_id: str) -> bool:
|
|
"""删除 Agent"""
|
|
agent = Agent.query.get(agent_id)
|
|
if not agent:
|
|
return False
|
|
|
|
# 更新 Gateway 的 Agent 计数
|
|
if agent.gateway_id:
|
|
gateway = Gateway.query.get(agent.gateway_id)
|
|
if gateway and gateway.agent_count > 0:
|
|
gateway.agent_count -= 1
|
|
|
|
db.session.delete(agent)
|
|
db.session.commit()
|
|
|
|
return True
|
|
|
|
def get_agent_stats(self) -> dict:
|
|
"""获取 Agent 统计"""
|
|
total = Agent.query.count()
|
|
online = Agent.query.filter_by(status='online').count()
|
|
offline = Agent.query.filter_by(status='offline').count()
|
|
busy = Agent.query.filter_by(status='busy').count()
|
|
|
|
# 总连接数
|
|
total_sessions = db.session.query(
|
|
db.func.sum(Agent.current_sessions)
|
|
).scalar() or 0
|
|
|
|
return {
|
|
'total': total,
|
|
'online': online,
|
|
'offline': offline,
|
|
'busy': busy,
|
|
'total_sessions': total_sessions,
|
|
}
|
|
|
|
def check_offline_agents(self, timeout: int = 120) -> List[Agent]:
|
|
"""检查超时下线的 Agent"""
|
|
from datetime import timedelta
|
|
|
|
threshold = datetime.utcnow() - timedelta(seconds=timeout)
|
|
|
|
offline_agents = Agent.query.filter(
|
|
Agent.status == 'online',
|
|
Agent.last_heartbeat < threshold
|
|
).all()
|
|
|
|
for agent in offline_agents:
|
|
agent.status = 'offline'
|
|
|
|
db.session.commit()
|
|
|
|
return offline_agents
|