more_dots/workflows/workflow_manager.py

"""
工作流管理器模块

支持动态注册和管理工作流类型
"""

from typing import Any, Callable, Dict, List, Optional, Type, Union
from enum import Enum
from datetime import datetime, timedelta
import logging

from agent.conversation import ConversationAgent
from agent.tool import ToolAgent
from core.registry import WorkflowRegistry, WorkflowMetadata

logger = logging.getLogger(__name__)


class WorkflowType(Enum):
    """可用的工作流类型"""
    CONVERSATION = "conversation"
    TOOL_USING = "tool_using"


class WorkflowManager:
    """
    管理不同工作流类型及其执行
    
    支持特性：
    - 动态注册工作流
    - 会话管理
    - 工作流元数据
    """
    
    def __init__(self, default_model_section: Optional[str] = None):
        self._default_model_section = default_model_section
        self._workflows: Dict[str, Any] = {}
        self._workflow_metadata: Dict[str, WorkflowMetadata] = {}
        self.active_sessions: Dict[str, Any] = {}
        
        self._register_default_workflows()
    
    def _register_default_workflows(self) -> None:
        """注册默认工作流"""
        self.register_workflow(
            name=WorkflowType.CONVERSATION.value,
            agent=ConversationAgent(model_section=self._default_model_section),
            description="多轮对话工作流",
            version="1.0.0",
        )
        self.register_workflow(
            name=WorkflowType.TOOL_USING.value,
            agent=ToolAgent(model_section=self._default_model_section),
            description="工具调用工作流",
            version="1.0.0",
        )
    
    def register_workflow(
        self,
        name: str,
        agent: Any,
        description: str = "",
        version: str = "1.0.0",
        default_model: Optional[str] = None,
        supported_features: Optional[List[str]] = None,
    ) -> None:
        """
        注册工作流
        
        Args:
            name: 工作流名称
            agent: Agent 实例
            description: 描述
            version: 版本
            default_model: 默认模型
            supported_features: 支持的特性列表
        """
        metadata = WorkflowMetadata(
            name=name,
            description=description,
            version=version,
            default_model=default_model,
            supported_features=supported_features or [],
        )
        
        self._workflows[name] = agent
        self._workflow_metadata[name] = metadata
        
        WorkflowRegistry._entries[name] = type(
            "RegistryEntry",
            (),
            {"instance": agent, "metadata": {"workflow_metadata": metadata}}
        )()
        
        logger.info(f"Registered workflow: {name} (v{version})")
    
    def unregister_workflow(self, name: str) -> bool:
        """
        注销工作流
        
        Args:
            name: 工作流名称
            
        Returns:
            是否成功注销
        """
        if name in self._workflows:
            del self._workflows[name]
            del self._workflow_metadata[name]
            WorkflowRegistry.unregister(name)
            logger.info(f"Unregistered workflow: {name}")
            return True
        return False
    
    def get_workflow(self, workflow_type: Union[WorkflowType, str]) -> Optional[Any]:
        """
        获取工作流实例
        
        Args:
            workflow_type: 工作流类型（枚举或字符串）
            
        Returns:
            Agent 实例
        """
        name = workflow_type.value if isinstance(workflow_type, WorkflowType) else workflow_type
        return self._workflows.get(name)
    
    def get_workflow_metadata(self, name: str) -> Optional[WorkflowMetadata]:
        """获取工作流元数据"""
        return self._workflow_metadata.get(name)
    
    def execute_workflow(
        self,
        workflow_type: Union[WorkflowType, str],
        user_input: str,
        session_id: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        执行指定工作流
        
        Args:
            workflow_type: 工作流类型
            user_input: 用户输入
            session_id: 会话ID
            **kwargs: 额外参数
            
        Returns:
            执行结果
        """
        name = workflow_type.value if isinstance(workflow_type, WorkflowType) else workflow_type
        workflow = self.get_workflow(name)
        
        if not workflow:
            return {"error": f"Workflow {name} not found"}
        
        if not session_id:
            session_id = f"session_{len(self.active_sessions) + 1}"
        
        result = workflow.run(user_input, **kwargs)
        
        self.active_sessions[session_id] = {
            "workflow_type": name,
            "last_result": result,
            "timestamp": self._get_timestamp()
        }
        
        return {
            "session_id": session_id,
            "workflow_type": name,
            "result": result
        }
    
    def get_available_workflows(self) -> List[str]:
        """获取可用工作流列表"""
        return list(self._workflows.keys())
    
    def get_workflow_info(self, name: str) -> Optional[Dict[str, Any]]:
        """获取工作流详细信息"""
        if name not in self._workflows:
            return None
        
        metadata = self._workflow_metadata.get(name)
        return {
            "name": name,
            "description": metadata.description if metadata else "",
            "version": metadata.version if metadata else "unknown",
            "supported_features": metadata.supported_features if metadata else [],
        }
    
    def _get_timestamp(self) -> str:
        """获取当前时间戳"""
        return datetime.now().isoformat()
    
    def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
        """获取会话信息"""
        return self.active_sessions.get(session_id)
    
    def cleanup_sessions(self, older_than_hours: int = 24) -> int:
        """
        清理过期会话
        
        Args:
            older_than_hours: 超过多少小时的会话将被清理
            
        Returns:
            清理的会话数量
        """
        cutoff_time = datetime.now() - timedelta(hours=older_than_hours)
        
        sessions_to_remove = []
        for session_id, session_data in self.active_sessions.items():
            session_time = datetime.fromisoformat(session_data["timestamp"])
            if session_time < cutoff_time:
                sessions_to_remove.append(session_id)
        
        for session_id in sessions_to_remove:
            del self.active_sessions[session_id]
        
        if sessions_to_remove:
            logger.info(f"Cleaned up {len(sessions_to_remove)} expired sessions")
        
        return len(sessions_to_remove)
    
    def create_agent_instance(
        self,
        workflow_name: str,
        agent_class: Type,
        model_section: Optional[str] = None,
        **kwargs
    ) -> Any:
        """
        创建并注册新的 Agent 实例
        
        Args:
            workflow_name: 工作流名称
            agent_class: Agent 类
            model_section: 模型配置段
            **kwargs: Agent 构造参数
            
        Returns:
            Agent 实例
        """
        model = model_section or self._default_model_section
        agent = agent_class(model_section=model, **kwargs)
        self.register_workflow(name=workflow_name, agent=agent)
        return agent