pit-router/app/services/agent_service.py

"""
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