feat: Phase 2 - 服务层实现 + Bug修复
This commit is contained in:
198
app/services/agent_service.py
Normal file
198
app/services/agent_service.py
Normal file
@@ -0,0 +1,198 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user