AIChatRoom/backend/models/agent.py

"""
Agent数据模型
定义AI聊天代理的配置结构
"""
from datetime import datetime
from typing import Optional, Dict, Any, List
from pydantic import Field
from beanie import Document


class AgentCapabilities:
    """Agent能力配置"""
    memory_enabled: bool = False          # 是否启用记忆
    mcp_tools: List[str] = []             # 可用的MCP工具
    skills: List[str] = []                # 可用的技能
    multimodal: bool = False              # 是否支持多模态
    
    def __init__(
        self,
        memory_enabled: bool = False,
        mcp_tools: Optional[List[str]] = None,
        skills: Optional[List[str]] = None,
        multimodal: bool = False
    ):
        self.memory_enabled = memory_enabled
        self.mcp_tools = mcp_tools or []
        self.skills = skills or []
        self.multimodal = multimodal
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "memory_enabled": self.memory_enabled,
            "mcp_tools": self.mcp_tools,
            "skills": self.skills,
            "multimodal": self.multimodal
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "AgentCapabilities":
        """从字典创建"""
        if not data:
            return cls()
        return cls(
            memory_enabled=data.get("memory_enabled", False),
            mcp_tools=data.get("mcp_tools", []),
            skills=data.get("skills", []),
            multimodal=data.get("multimodal", False)
        )


class AgentBehavior:
    """Agent行为配置"""
    speak_threshold: float = 0.5          # 发言阈值(判断是否需要发言)
    max_speak_per_round: int = 2          # 每轮最多发言次数
    speak_style: str = "balanced"         # 发言风格: concise, balanced, detailed
    
    def __init__(
        self,
        speak_threshold: float = 0.5,
        max_speak_per_round: int = 2,
        speak_style: str = "balanced"
    ):
        self.speak_threshold = speak_threshold
        self.max_speak_per_round = max_speak_per_round
        self.speak_style = speak_style
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "speak_threshold": self.speak_threshold,
            "max_speak_per_round": self.max_speak_per_round,
            "speak_style": self.speak_style
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "AgentBehavior":
        """从字典创建"""
        if not data:
            return cls()
        return cls(
            speak_threshold=data.get("speak_threshold", 0.5),
            max_speak_per_round=data.get("max_speak_per_round", 2),
            speak_style=data.get("speak_style", "balanced")
        )


class Agent(Document):
    """
    Agent文档模型
    存储AI代理的配置信息
    """
    agent_id: str = Field(..., description="唯一标识")
    name: str = Field(..., description="Agent名称")
    role: str = Field(..., description="角色定义")
    system_prompt: str = Field(..., description="系统提示词")
    provider_id: str = Field(..., description="使用的AI接口ID")
    
    # 模型参数
    temperature: float = Field(default=0.7, ge=0, le=2, description="温度参数")
    max_tokens: int = Field(default=2000, gt=0, description="最大token数")
    
    # 能力配置
    capabilities: Dict[str, Any] = Field(
        default_factory=lambda: {
            "memory_enabled": False,
            "mcp_tools": [],
            "skills": [],
            "multimodal": False
        },
        description="能力配置"
    )
    
    # 行为配置
    behavior: Dict[str, Any] = Field(
        default_factory=lambda: {
            "speak_threshold": 0.5,
            "max_speak_per_round": 2,
            "speak_style": "balanced"
        },
        description="行为配置"
    )
    
    # 外观配置
    avatar: Optional[str] = Field(default=None, description="头像URL")
    color: str = Field(default="#1890ff", description="代表颜色")
    
    # 元数据
    enabled: bool = Field(default=True, description="是否启用")
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    
    class Settings:
        name = "agents"
    
    def get_capabilities(self) -> AgentCapabilities:
        """获取能力配置对象"""
        return AgentCapabilities.from_dict(self.capabilities)
    
    def get_behavior(self) -> AgentBehavior:
        """获取行为配置对象"""
        return AgentBehavior.from_dict(self.behavior)
    
    class Config:
        json_schema_extra = {
            "example": {
                "agent_id": "product-manager",
                "name": "产品经理",
                "role": "产品规划和需求分析专家",
                "system_prompt": "你是一位经验丰富的产品经理，擅长分析用户需求...",
                "provider_id": "openrouter-gpt4",
                "temperature": 0.7,
                "max_tokens": 2000,
                "capabilities": {
                    "memory_enabled": True,
                    "mcp_tools": ["web_search"],
                    "skills": [],
                    "multimodal": False
                },
                "behavior": {
                    "speak_threshold": 0.5,
                    "max_speak_per_round": 2,
                    "speak_style": "balanced"
                },
                "avatar": "https://example.com/avatar.png",
                "color": "#1890ff"
            }
        }