MemoryBear/api/app/services/multi_agent_service.py

"""多 Agent 配置管理服务"""
import uuid
import json
from typing import Optional, List, Tuple, Any, Annotated

from fastapi import Depends
from sqlalchemy.orm import Session
from sqlalchemy import select, desc

from app.db import get_db
from app.models import MultiAgentConfig, App, AgentConfig
from app.schemas.multi_agent_schema import (
    MultiAgentConfigCreate,
    MultiAgentConfigUpdate,
    MultiAgentRunRequest
)
from app.services.model_service import ModelApiKeyService
from app.services.multi_agent_orchestrator import MultiAgentOrchestrator
from app.core.exceptions import ResourceNotFoundException, BusinessException
from app.core.error_codes import BizCode
from app.core.logging_config import get_business_logger
from app.models import AppRelease

logger = get_business_logger()


def convert_uuids_to_str(obj: Any) -> Any:
    """递归转换对象中的所有 UUID 为字符串

    Args:
        obj: 要转换的对象（dict, list, UUID 等）

    Returns:
        转换后的对象
    """
    if isinstance(obj, uuid.UUID):
        return str(obj)
    elif isinstance(obj, dict):
        return {k: convert_uuids_to_str(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_uuids_to_str(item) for item in obj]
    else:
        return obj


class MultiAgentService:
    """多 Agent 配置管理服务"""

    def __init__(self, db: Session):
        self.db = db

    def create_config(
        self,
        app_id: uuid.UUID,
        data: MultiAgentConfigCreate,
        created_by: uuid.UUID
    ) -> MultiAgentConfig:
        """创建多 Agent 配置

        Args:
            app_id: 应用 ID
            data: 配置数据
            created_by: 创建者 ID

        Returns:
            多 Agent 配置
        """
        # 1. 验证应用存在
        app = self.db.get(App, app_id)
        if not app:
            raise ResourceNotFoundException("应用", str(app_id))

        # 2. 检查是否已有有效配置
        existing = self.db.scalars(
            select(MultiAgentConfig)
            .where(
                MultiAgentConfig.app_id == app_id,
                MultiAgentConfig.is_active.is_(True)
            )
            .order_by(MultiAgentConfig.updated_at.desc())
        ).first()
        if existing:
            raise BusinessException("应用已有多 Agent 配置", BizCode.DUPLICATE_RESOURCE)

        # 3. 验证主 Agent 存在
        master_agent = self.db.get(AgentConfig, data.master_agent_id)
        if not master_agent:
            raise ResourceNotFoundException("主 Agent", str(data.master_agent_id))

        # 4. 验证子 Agent 存在
        for sub_agent in data.sub_agents:
            agent = self.db.get(AgentConfig, sub_agent.agent_id)
            if not agent:
                raise ResourceNotFoundException("子 Agent", str(sub_agent.agent_id))

        # 5. 创建配置（转换 UUID 为字符串以支持 JSON 序列化）
        sub_agents_data = [convert_uuids_to_str(sub_agent.model_dump()) for sub_agent in data.sub_agents]
        routing_rules_data = [convert_uuids_to_str(rule.model_dump()) for rule in data.routing_rules] if data.routing_rules else None

        # 处理 execution_config（可能是 None、字典或 Pydantic 模型）
        if data.execution_config is None:
            execution_config_data = {}
        elif isinstance(data.execution_config, dict):
            execution_config_data = convert_uuids_to_str(data.execution_config)
        else:
            execution_config_data = convert_uuids_to_str(data.execution_config.model_dump())

        config = MultiAgentConfig(
            app_id=app_id,
            master_agent_id=data.master_agent_id,
            master_agent_name=data.master_agent_name,
            orchestration_mode=data.orchestration_mode,
            sub_agents=sub_agents_data,
            routing_rules=routing_rules_data,
            execution_config=execution_config_data,
            aggregation_strategy=data.aggregation_strategy
        )

        self.db.add(config)
        self.db.commit()
        self.db.refresh(config)

        logger.info(
            "创建多 Agent 配置成功",
            extra={
                "config_id": str(config.id),
                "app_id": str(app_id),
                "mode": data.orchestration_mode,
                "sub_agent_count": len(data.sub_agents)
            }
        )

        return config

    def get_config(self, app_id: uuid.UUID) -> Optional[MultiAgentConfig]:
        """获取多 Agent 配置

        Args:
            app_id: 应用 ID

        Returns:
            多 Agent 配置，如果不存在返回 None
        """
        return self.db.scalars(
            select(MultiAgentConfig)
            .where(
                MultiAgentConfig.app_id == app_id,
                MultiAgentConfig.is_active.is_(True)
            )
            .order_by(MultiAgentConfig.updated_at.desc())
        ).first()

    def get_multi_agent_configs(self, app_id: uuid.UUID) -> Optional[dict]:
        """通过 app_id 获取最新有效的多智能体配置，并将 agent_id 转换为 app_id

        Args:
            app_id: 应用 ID

        Returns:
            转换后的配置字典，如果不存在返回 None
        """
        config = self.get_config(app_id)
        if not config:
            return None

        #兼容代码
        if not config.default_model_config_id:
            master_release = self.db.get(AppRelease, config.master_agent_id)
            config.default_model_config_id = master_release.default_model_config_id if master_release else None

        # 转换 sub_agents 中的 agent_id (release_id) 为 app_id
        converted_sub_agents = []
        for sub_agent in config.sub_agents:
            sub_agent_copy = sub_agent.copy()
            release_id = sub_agent.get("agent_id")
            if release_id:
                try:
                    release_id_uuid = uuid.UUID(release_id) if isinstance(release_id, str) else release_id
                    sub_release = self.db.get(AppRelease, release_id_uuid)
                    if sub_release:
                        sub_agent_copy["agent_id"] = str(sub_release.app_id)
                except Exception as e:
                    logger.warning(f"转换 sub_agent agent_id 失败: {release_id}, 错误: {str(e)}")
            converted_sub_agents.append(sub_agent_copy)

        # 构建返回的配置字典
        return {
            "id": config.id,
            "app_id": config.app_id,
            "default_model_config_id": config.default_model_config_id,
            "model_parameters": config.model_parameters,
            "orchestration_mode": config.orchestration_mode,
            "sub_agents": converted_sub_agents,
            "routing_rules": config.routing_rules,
            "execution_config": config.execution_config,
            "aggregation_strategy": config.aggregation_strategy,
            "is_active": config.is_active,
            "created_at": config.created_at,
            "updated_at": config.updated_at
        }

    def get_published_config_by_agent_id(self, agent_id: uuid.UUID) -> Optional[dict]:
        """通过 agent_id 获取当前发布版本的完整配置

        Args:
            agent_id: Agent 配置 ID

        Returns:
            当前发布版本的配置字典，如果没有发布版本则返回 None
        """
        from app.models import AppRelease

        # 查询 Agent 配置
        agent_config = self.db.get(AgentConfig, agent_id)
        if not agent_config:
            logger.warning(f"Agent 配置不存在: {agent_id}")
            return None

        # 获取关联的应用
        app = self.db.get(App, agent_config.app_id)
        if not app or not app.current_release_id:
            logger.warning(f"应用未发布或不存在: app_id={agent_config.app_id}")
            return None

        # 获取当前发布版本
        release = self.db.get(AppRelease, app.current_release_id)
        if not release:
            logger.warning(f"发布版本不存在: release_id={app.current_release_id}")
            return None

        # 从发布版本的 config 中获取完整配置
        # config 是一个 JSON 对象，包含了发布时的配置快照
        config_data = release.config
        if config_data and isinstance(config_data, dict):
            return config_data

        return None

    def get_published_by_agent_id(self, agent_id: uuid.UUID) -> Optional[AppRelease]:
        """通过 agent_id 获取当前发布版本的完整配置

        Args:
            agent_id: Agent 配置 ID

        Returns:
            当前发布版本的配置字典，如果没有发布版本则返回 None
        """

        # 获取关联的应用
        app = self.db.get(App, agent_id)
        if not app or not app.current_release_id:
            logger.warning(f"应用未发布或不存在: app_id={agent_id}")
            return None

        # 获取当前发布版本
        release = self.db.get(AppRelease, app.current_release_id)
        if not release:
            logger.warning(f"发布版本不存在: release_id={app.current_release_id}")
            return None
        return release

    def check_config_data(self,app_id: uuid.UUID, data: MultiAgentConfigUpdate) -> MultiAgentConfig:
        # 1. 验证应用存在
        app = self.db.get(App, app_id)
        if not app:
            raise ResourceNotFoundException("应用", str(app_id))

        # 2. 验证模型配置（如果提供了）
        if data.default_model_config_id:
            model_api_key = ModelApiKeyService.get_a_api_key(self.db, data.default_model_config_id)
            if not model_api_key:
                raise ResourceNotFoundException("模型配置", str(data.default_model_config_id))

        # 3. 验证子 Agent 存在并获取发布版本 ID
        for sub_agent in data.sub_agents:
            agent_app_release = self.get_published_by_agent_id(sub_agent.agent_id)
            if not agent_app_release:
                raise ResourceNotFoundException("子 Agent 未发布或不存在", str(sub_agent.agent_id))

                # 使用发布版本 ID
            sub_agent.agent_id = agent_app_release.id

        # 5. 创建配置（转换 UUID 为字符串以支持 JSON 序列化）
        sub_agents_data = [convert_uuids_to_str(sub_agent.model_dump()) for sub_agent in data.sub_agents]
        # routing_rules_data = [convert_uuids_to_str(rule.model_dump()) for rule in data.routing_rules] if data.routing_rules else None

        # 处理 execution_config（可能是 None、字典或 Pydantic 模型）
        if data.execution_config is None:
            execution_config_data = {}
        elif isinstance(data.execution_config, dict):
                execution_config_data = convert_uuids_to_str(data.execution_config)
        else:
            execution_config_data = convert_uuids_to_str(data.execution_config.model_dump())

        # 处理 model_parameters（可能是 None、字典或 Pydantic 模型）
        if data.model_parameters is None:
            model_parameters_data = None
        # elif isinstance(data.model_parameters, dict):
        #     # 过滤掉值为 None 的字段
        #     model_parameters_data = {k: v for k, v in data.model_parameters.items() if v is not None}
        else:
            # 过滤掉值为 None 的字段
            # model_parameters_data = {k: v for k, v in data.model_parameters.model_dump().items() if v is not None}
            model_parameters_data = data.model_parameters

        config = MultiAgentConfig(
                app_id=app_id,
                master_agent_id=data.master_agent_id,
                master_agent_name=data.master_agent_name,
                default_model_config_id=data.default_model_config_id,
                model_parameters=model_parameters_data,
                orchestration_mode=data.orchestration_mode,
                sub_agents=sub_agents_data,
                # routing_rules=routing_rules_data,
                execution_config=execution_config_data,
                aggregation_strategy=data.aggregation_strategy
            )
        return config

    def update_config(
        self,
        app_id: uuid.UUID,
        data: MultiAgentConfigUpdate
    ) -> MultiAgentConfig:
        """更新多 Agent 配置

        Args:
            app_id: 应用 ID
            data: 更新数据

        Returns:
            更新后的配置
        """
        config = self.get_config(app_id)
        newConfig = self.check_config_data(app_id, data)
        if not config:
            config = newConfig
            self.db.add(config)
            self.db.commit()
            self.db.refresh(config)
            logger.info(
                "创建多 Agent 配置成功",
                extra={
                    "config_id": str(config.id),
                    "app_id": str(app_id),
                    "mode": data.orchestration_mode,
                    "sub_agent_count": len(data.sub_agents)
                }
            )
            return config

        # 完全替换配置，但对于数据库 NOT NULL 字段，如果新值是 None 则保留原值
        config.default_model_config_id = newConfig.default_model_config_id
        config.model_parameters = newConfig.model_parameters
        config.orchestration_mode = newConfig.orchestration_mode or config.orchestration_mode
        config.sub_agents = newConfig.sub_agents if newConfig.sub_agents is not None else config.sub_agents
        config.routing_rules = newConfig.routing_rules
        config.execution_config = newConfig.execution_config if newConfig.execution_config else config.execution_config
        config.aggregation_strategy = newConfig.aggregation_strategy or config.aggregation_strategy
        self.db.commit()
        self.db.refresh(config)

        logger.info(
            "更新多 Agent 配置成功",
            extra={
                "config_id": str(config.id),
                "app_id": str(app_id)
            }
        )

        return config

    def delete_config(self, app_id: uuid.UUID) -> None:
        """删除多 Agent 配置

        Args:
            app_id: 应用 ID
        """
        config = self.get_config(app_id)
        if not config:
            raise ResourceNotFoundException("多 Agent 配置", str(app_id))

        # 逻辑删除多 Agent 配置
        config.is_active = False
        self.db.commit()

        logger.info(
            "删除多 Agent 配置成功",
            extra={
                "config_id": str(config.id),
                "app_id": str(app_id)
            }
        )

    async def run(
        self,
        app_id: uuid.UUID,
        request: MultiAgentRunRequest
    ) -> dict:
        """运行多 Agent 任务

        Args:
            app_id: 应用 ID
            request: 运行请求

        Returns:
            执行结果
        """
        # 1. 获取配置
        config = self.get_config(app_id)
        if not config:
            raise ResourceNotFoundException("多 Agent 配置", str(app_id))

        if not config.is_active:
            raise BusinessException("多 Agent 配置已禁用", BizCode.RESOURCE_DISABLED)

        # 2. 创建编排器
        orchestrator = MultiAgentOrchestrator(self.db, config)

        # 3. 执行任务
        result = await orchestrator.execute(
            message=request.message,
            conversation_id=request.conversation_id,
            user_id=request.user_id,
            variables=request.variables,
            use_llm_routing=getattr(request, 'use_llm_routing', True),  # 默认启用 LLM 路由
            web_search=getattr(request, 'web_search', False),  # 网络搜索参数
            memory=getattr(request, 'memory', True)  # 记忆功能参数
        )

        await self._save_conversation_message(
            conversation_id=request.conversation_id,
            user_message=request.message,
            assistant_message=result.get("message", ""),
            app_id=app_id,
            user_id=request.user_id,
            meta_data={
                "mode": result.get("mode"),
                "elapsed_time": result.get("elapsed_time"),
                "usage": result.get("usage", {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                })
            }
        )

        return result

    async def run_stream(
        self,
        app_id: uuid.UUID,
        request: MultiAgentRunRequest,
        storage_type :str,
        user_rag_memory_id :str
    ):
        """运行多 Agent 任务（流式返回）

        Args:
            app_id: 应用 ID
            request: 运行请求

        Yields:
            SSE 格式的事件流
        """
        # 1. 获取配置
        config = self.get_config(app_id)
        if not config:
            raise ResourceNotFoundException("多 Agent 配置", str(app_id))

        if not config.is_active:
            raise BusinessException("多 Agent 配置已禁用", BizCode.NOT_FOUND)

        # 2. 创建编排器
        orchestrator = MultiAgentOrchestrator(self.db, config)

        full_content = ""
        total_tokens = 0

        # 3. 流式执行任务
        async for event in orchestrator.execute_stream(
            message=request.message,
            conversation_id=request.conversation_id,
            user_id=request.user_id,
            variables=request.variables,
            use_llm_routing=getattr(request, 'use_llm_routing', True),
            web_search=getattr(request, 'web_search', False),  # 网络搜索参数
            memory=getattr(request, 'memory', True) , # 记忆功能参数
            storage_type=storage_type,
            user_rag_memory_id=user_rag_memory_id
        ):
            if "sub_usage" in event:
                if "data:" in event:
                    try:
                        data_line = event.split("data: ", 1)[1].strip()
                        data = json.loads(data_line)
                        if "total_tokens" in data:
                            total_tokens += data["total_tokens"]
                    except:
                        pass
            else:
                yield event
                if "data:" in event:
                    try:
                        data_line = event.split("data: ", 1)[1].strip()
                        data = json.loads(data_line)
                        if "content" in data:
                            full_content += data["content"]
                    except:
                        pass

        await self._save_conversation_message(
            conversation_id=request.conversation_id,
            user_message=request.message,
            assistant_message=full_content,
            app_id=app_id,
            user_id=request.user_id,
            meta_data={
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": total_tokens
                }
            }
        )

    async def _save_conversation_message(
        self,
        conversation_id: uuid.UUID,
        user_message: str,
        assistant_message: str,
        meta_data: dict,
        app_id: Optional[uuid.UUID] = None,
        user_id: Optional[str] = None
    ) -> None:
        """保存会话消息

        Args:
            conversation_id: 会话ID
            user_message: 用户消息
            assistant_message: AI 回复消息
            meta_data: 元数据（包括 token 消耗）
            app_id: 应用ID
            user_id: 用户ID
        """
        try:
            from app.services.conversation_service import ConversationService

            conversation_service = ConversationService(self.db)

            conversation_service.add_message(
                conversation_id=conversation_id,
                role="user",
                content=user_message
            )
            conversation_service.add_message(
                conversation_id=conversation_id,
                role="assistant",
                content=assistant_message,
                meta_data=meta_data
            )

            logger.debug(
                "保存多 Agent 会话消息",
                extra={
                    "conversation_id": conversation_id,
                    "user_message_length": len(user_message),
                    "assistant_message_length": len(assistant_message)
                }
            )

        except Exception as e:
            logger.warning("保存会话消息失败", extra={"error": str(e)})

    # def add_sub_agent(
    #     self,
    #     app_id: uuid.UUID,
    #     agent_id: uuid.UUID,
    #     name: str,
    #     role: Optional[str] = None,
    #     priority: int = 1,
    #     capabilities: Optional[List[str]] = None
    # ) -> MultiAgentConfig:
    #     """添加子 Agent

    #     Args:
    #         app_id: 应用 ID
    #         agent_id: Agent ID
    #         name: Agent 名称
    #         role: 角色描述
    #         priority: 优先级
    #         capabilities: 能力列表

    #     Returns:
    #         更新后的配置
    #     """
    #     config = self.get_config(app_id)
    #     if not config:
    #         raise ResourceNotFoundException("多 Agent 配置", str(app_id))

    #     # 验证 Agent 存在
    #     agent = self.db.get(AgentConfig, agent_id)
    #     if not agent:
    #         raise ResourceNotFoundException("Agent", str(agent_id))

    #     # 检查是否已存在
    #     for sub_agent in config.sub_agents:
    #         if sub_agent["agent_id"] == str(agent_id):
    #             raise BusinessException("Agent 已存在于配置中", BizCode.DUPLICATE_RESOURCE)

    #     # 添加子 Agent
    #     new_sub_agent = {
    #         "agent_id": str(agent_id),
    #         "name": name,
    #         "role": role,
    #         "priority": priority,
    #         "capabilities": capabilities or []
    #     }

    #     config.sub_agents.append(new_sub_agent)

    #     # 标记为已修改
    #     self.db.add(config)
    #     self.db.commit()
    #     self.db.refresh(config)

    #     logger.info(
    #         "添加子 Agent 成功",
    #         extra={
    #             "config_id": str(config.id),
    #             "agent_id": str(agent_id),
    #             "agent_name": name
    #         }
    #     )

    #     return config

    # def remove_sub_agent(
    #     self,
    #     app_id: uuid.UUID,
    #     agent_id: uuid.UUID
    # ) -> MultiAgentConfig:
    #     """移除子 Agent

    #     Args:
    #         app_id: 应用 ID
    #         agent_id: Agent ID

    #     Returns:
    #         更新后的配置
    #     """
    #     config = self.get_config(app_id)
    #     if not config:
    #         raise ResourceNotFoundException("多 Agent 配置", str(app_id))

    #     # 查找并移除
    #     original_count = len(config.sub_agents)
    #     config.sub_agents = [
    #         sub_agent for sub_agent in config.sub_agents
    #         if sub_agent["agent_id"] != str(agent_id)
    #     ]

    #     if len(config.sub_agents) == original_count:
    #         raise ResourceNotFoundException("子 Agent", str(agent_id))

    #     # 标记为已修改
    #     self.db.add(config)
    #     self.db.commit()
    #     self.db.refresh(config)

    #     logger.info(
    #         "移除子 Agent 成功",
    #         extra={
    #             "config_id": str(config.id),
    #             "agent_id": str(agent_id)
    #         }
    #     )

    #     return config

    def list_configs(
        self,
        workspace_id: uuid.UUID,
        page: int = 1,
        pagesize: int = 20
    ) -> Tuple[List[MultiAgentConfig], int]:
        """列出多 Agent 配置

        Args:
            workspace_id: 工作空间 ID
            page: 页码
            pagesize: 每页数量

        Returns:
            配置列表和总数
        """
        # 构建查询
        stmt = (
            select(MultiAgentConfig)
            .join(App)
            .where(App.workspace_id == workspace_id)
            .order_by(desc(MultiAgentConfig.created_at))
        )

        # 总数
        count_stmt = stmt.with_only_columns(MultiAgentConfig.id)
        total = len(self.db.execute(count_stmt).all())

        # 分页
        stmt = stmt.offset((page - 1) * pagesize).limit(pagesize)
        configs = list(self.db.scalars(stmt).all())

        return configs, total

# ==================== 依赖注入函数 ====================

def get_multi_agent_service(
        db: Annotated[Session, Depends(get_db)]
) -> MultiAgentService:
    """获取工作流服务（依赖注入）"""
    return MultiAgentService(db)