🎯 Multi-Agent 技能系统

"""
基础技能定义 - 使用 @tool 装饰器
"""
from langchain.tools import tool
from typing import List, Dict

# ========== 搜索类技能 ==========

@tool
def search_internet(query: str, max_results: int = 5) -> List[Dict]:
    """搜索互联网信息

    Args:
        query: 搜索关键词
        max_results: 最大结果数量

    Returns:
        搜索结果列表
    """
    # 实际会调用 Tavily 或其他搜索 API
    return [
        {"title": "结果1", "url": "https://...", "snippet": "摘要..."},
        {"title": "结果2", "url": "https://...", "snippet": "摘要..."},
    ]

@tool
def search_database(query: str, table: str) -> List[Dict]:
    """搜索数据库

    Args:
        query: SQL 查询或自然语言查询
        table: 目标表名

    Returns:
        查询结果
    """
    # 实际会连接数据库执行查询
    return [{"id": 1, "name": "数据1"}, {"id": 2, "name": "数据2"}]

# ========== 数据处理类技能 ==========

@tool
def analyze_data(data: str, analysis_type: str = "summary") -> Dict:
    """分析数据并生成洞察

    Args:
        data: 待分析的数据
        analysis_type: 分析类型（summary, trend, anomaly）

    Returns:
        分析结果
    """
    return {
        "insights": ["洞察1", "洞察2"],
        "summary": "数据摘要...",
        "score": 0.85
    }

@tool
def transform_data(data: str, format: str) -> str:
    """转换数据格式

    Args:
        data: 原始数据
        format: 目标格式（json, csv, xml）

    Returns:
        转换后的数据
    """
    # 实际会执行格式转换
    return f"转换为 {format} 格式的数据"

# ========== 通信类技能 ==========

@tool
def send_email(to: str, subject: str, body: str) -> bool:
    """发送邮件

    Args:
        to: 收件人
        subject: 邮件主题
        body: 邮件正文

    Returns:
        是否发送成功
    """
    # 实际会调用邮件服务
    print(f"发送邮件给 {to}: {subject}")
    return True

@tool
def send_notification(user_id: str, message: str, priority: str = "normal") -> bool:
    """发送通知

    Args:
        user_id: 用户ID
        message: 通知内容
        priority: 优先级（low, normal, high）

    Returns:
        是否发送成功
    """
    print(f"通知用户 {user_id} ({priority}): {message}")
    return True

# ========== 存储类技能 ==========

@tool
def save_to_database(data: Dict, table: str) -> bool:
    """保存数据到数据库

    Args:
        data: 要保存的数据
        table: 目标表名

    Returns:
        是否保存成功
    """
    print(f"保存到表 {table}: {data}")
    return True

@tool
def save_file(content: str, filepath: str) -> bool:
    """保存文件

    Args:
        content: 文件内容
        filepath: 文件路径

    Returns:
        是否保存成功
    """
    with open(filepath, 'w') as f:
        f.write(content)
    return True

# ========== 格式化类技能 ==========

@tool
def format_report(data: Dict, template: str = "markdown") -> str:
    """格式化报告

    Args:
        data: 报告数据
        template: 报告模板（markdown, html, pdf）

    Returns:
        格式化后的报告
    """
    # 实际会使用模板引擎
    return f"""
# 报告
## 数据摘要
{data.get('summary', 'N/A')}

## 详细内容
{data.get('details', 'N/A')}
"""

# 优势：
# - 清晰的技能分类
# - 完整的类型注解
# - 详细的文档字符串
# - 易于理解和使用

"""
技能组织和管理 - 技能库模式
"""
from typing import List
from langchain.tools import BaseTool

class SkillRegistry:
    """技能注册表：集中管理所有技能"""

    def __init__(self):
        self._skills = {}
        self._categories = {
            "search": [],
            "data": [],
            "communication": [],
            "storage": [],
            "formatting": []
        }

    def register(self, skill: BaseTool, category: str):
        """注册技能"""
        self._skills[skill.name] = skill
        if category in self._categories:
            self._categories[category].append(skill)

    def get_skill(self, name: str) -> BaseTool:
        """获取单个技能"""
        return self._skills.get(name)

    def get_category(self, category: str) -> List[BaseTool]:
        """获取某类别的所有技能"""
        return self._categories.get(category, [])

    def get_skills(self, names: List[str]) -> List[BaseTool]:
        """批量获取技能"""
        return [self._skills[name] for name in names if name in self._skills]

    def list_all(self) -> List[str]:
        """列出所有技能名称"""
        return list(self._skills.keys())

# 创建全局技能注册表
skill_registry = SkillRegistry()

# 注册技能
skill_registry.register(search_internet, "search")
skill_registry.register(search_database, "search")
skill_registry.register(analyze_data, "data")
skill_registry.register(transform_data, "data")
skill_registry.register(send_email, "communication")
skill_registry.register(send_notification, "communication")
skill_registry.register(save_to_database, "storage")
skill_registry.register(save_file, "storage")
skill_registry.register(format_report, "formatting")

# 使用示例：为代理配置技能
research_skills = skill_registry.get_category("search") + \
                  skill_registry.get_category("data")

communication_skills = skill_registry.get_category("communication") + \
                       skill_registry.get_category("storage")

# 优势：
# - 集中管理，易于维护
# - 分类组织，清晰结构
# - 灵活组合，按需获取
# - 可扩展，易于添加新技能

"""
模式 1：完全共享 - 所有代理使用相同技能
"""
from langgraph.graph import StateGraph, MessagesState
from langchain_openai import ChatOpenAI

# 定义共享技能集
shared_skills = [
    search_internet,
    analyze_data,
    send_notification,
    format_report
]

# 代理 1：研究专员
def research_agent(state: MessagesState):
    """研究专员：使用共享技能"""
    llm = ChatOpenAI(model="gpt-4o")
    llm_with_tools = llm.bind_tools(shared_skills)  # 使用共享技能
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# 代理 2：分析专员
def analysis_agent(state: MessagesState):
    """分析专员：使用相同的共享技能"""
    llm = ChatOpenAI(model="claude-sonnet-4-5")
    llm_with_tools = llm.bind_tools(shared_skills)  # 使用相同技能
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# 代理 3：报告专员
def report_agent(state: MessagesState):
    """报告专员：使用相同的共享技能"""
    llm = ChatOpenAI(model="gpt-4o")
    llm_with_tools = llm.bind_tools(shared_skills)  # 使用相同技能
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# 构建工作流
graph = StateGraph(MessagesState)
graph.add_node("research", research_agent)
graph.add_node("analysis", analysis_agent)
graph.add_node("report", report_agent)

# 优势：
# - 所有代理能力一致
# - 便于协作和交接
# - 维护简单（更新一处即可）
# - 适合：能力需求相似的代理团队

"""
模式 2：分层共享 - 基础技能 + 专用技能
"""
# 基础技能：所有代理都有
base_skills = [
    search_internet,
    send_notification
]

# 高级技能：仅部分代理使用
advanced_data_skills = [
    analyze_data,
    transform_data,
    save_to_database
]

advanced_comm_skills = [
    send_email,
    format_report
]

# L1 代理：仅基础技能
def l1_agent(state: MessagesState):
    """L1 代理：基础能力"""
    llm = ChatOpenAI(model="gpt-4o-mini")
    llm_with_tools = llm.bind_tools(base_skills)  # 仅基础技能
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

# L2 代理：基础 + 数据处理
def l2_data_agent(state: MessagesState):
    """L2 数据代理：基础 + 高级数据技能"""
    llm = ChatOpenAI(model="gpt-4o")
    combined_skills = base_skills + advanced_data_skills
    llm_with_tools = llm.bind_tools(combined_skills)
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

# L2 代理：基础 + 通信
def l2_comm_agent(state: MessagesState):
    """L2 通信代理：基础 + 高级通信技能"""
    llm = ChatOpenAI(model="gpt-4o")
    combined_skills = base_skills + advanced_comm_skills
    llm_with_tools = llm.bind_tools(combined_skills)
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

# 优势：
# - 权限分级（基础 vs 高级）
# - 专业化分工（数据 vs 通信）
# - 成本优化（基础代理用便宜模型）
# - 适合：有明确层级和专业分工的系统

"""
模式 3：技能组合 - 预定义技能组
"""
# 定义技能组
skill_groups = {
    "researcher": [
        search_internet,
        search_database,
        analyze_data
    ],
    "writer": [
        analyze_data,
        format_report,
        save_file
    ],
    "communicator": [
        send_email,
        send_notification,
        save_to_database
    ],
    "analyst": [
        search_database,
        analyze_data,
        transform_data,
        format_report
    ]
}

def create_agent_with_skills(role: str, model: str = "gpt-4o"):
    """工厂函数：根据角色创建代理"""
    skills = skill_groups.get(role, [])

    def agent(state: MessagesState):
        llm = ChatOpenAI(model=model)
        llm_with_tools = llm.bind_tools(skills)
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

    return agent

# 创建不同角色的代理
researcher = create_agent_with_skills("researcher", "gpt-4o")
writer = create_agent_with_skills("writer", "claude-sonnet-4-5")
communicator = create_agent_with_skills("communicator", "gpt-4o")
analyst = create_agent_with_skills("analyst", "claude-sonnet-4-5")

# 优势：
# - 角色清晰（研究、写作、通信、分析）
# - 技能精准匹配
# - 易于扩展（添加新角色）
# - 适合：明确角色分工的团队

"""
完整实战示例 - 内容创作平台（技能系统）
"""
import os
from typing import TypedDict, Annotated, List
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain.tools import tool

# ========== 1. 定义技能库 ==========

# 研究类技能
@tool
def web_research(topic: str, depth: str = "normal") -> dict:
    """网络研究技能

    Args:
        topic: 研究主题
        depth: 深度（quick, normal, deep）
    """
    # 实际会调用搜索 API
    return {
        "sources": ["来源1", "来源2", "来源3"],
        "key_points": ["要点1", "要点2", "要点3"],
        "summary": f"关于 {topic} 的研究摘要..."
    }

@tool
def competitor_analysis(topic: str) -> dict:
    """竞品分析技能"""
    return {
        "competitors": ["竞品A", "竞品B"],
        "gaps": ["差异点1", "差异点2"],
        "opportunities": ["机会1", "机会2"]
    }

# 创作类技能
@tool
def generate_outline(topic: str, target_length: int) -> dict:
    """生成大纲技能

    Args:
        topic: 主题
        target_length: 目标字数
    """
    return {
        "title": f"{topic} - 完整指南",
        "sections": [
            {"title": "引言", "points": ["背景", "目标"]},
            {"title": "正文", "points": ["要点1", "要点2", "要点3"]},
            {"title": "结论", "points": ["总结", "展望"]}
        ]
    }

@tool
def check_grammar(text: str) -> dict:
    """语法检查技能"""
    return {
        "errors": [],
        "suggestions": ["建议1", "建议2"],
        "score": 95
    }

# SEO 类技能
@tool
def analyze_seo(content: str, keywords: List[str]) -> dict:
    """SEO 分析技能

    Args:
        content: 内容
        keywords: 目标关键词列表
    """
    return {
        "keyword_density": {"关键词1": 2.5, "关键词2": 1.8},
        "readability_score": 85,
        "recommendations": ["增加内链", "优化标题"]
    }

@tool
def generate_meta_tags(content: str) -> dict:
    """生成元标签技能"""
    return {
        "title": "SEO 优化的标题",
        "description": "吸引人的描述（< 160字）",
        "keywords": ["关键词1", "关键词2", "关键词3"]
    }

# 发布类技能
@tool
def publish_to_cms(content: str, metadata: dict) -> bool:
    """发布到 CMS 技能

    Args:
        content: 文章内容
        metadata: 元数据（标题、标签等）
    """
    print(f"发布文章：{metadata.get('title')}")
    return True

@tool
def schedule_social_media(post: str, platforms: List[str]) -> dict:
    """社交媒体排期技能

    Args:
        post: 社交媒体文案
        platforms: 目标平台列表
    """
    return {
        "scheduled": True,
        "platforms": platforms,
        "publish_time": "2024-01-15 10:00"
    }

# ========== 2. 组织技能库 ==========

class ContentSkillRegistry:
    """内容创作技能注册表"""

    RESEARCH_SKILLS = [web_research, competitor_analysis]
    WRITING_SKILLS = [generate_outline, check_grammar]
    SEO_SKILLS = [analyze_seo, generate_meta_tags]
    PUBLISHING_SKILLS = [publish_to_cms, schedule_social_media]

    @classmethod
    def get_researcher_skills(cls):
        """研究专员技能"""
        return cls.RESEARCH_SKILLS

    @classmethod
    def get_writer_skills(cls):
        """写作专员技能"""
        return cls.RESEARCH_SKILLS + cls.WRITING_SKILLS

    @classmethod
    def get_seo_specialist_skills(cls):
        """SEO 专员技能"""
        return cls.WRITING_SKILLS + cls.SEO_SKILLS

    @classmethod
    def get_publisher_skills(cls):
        """发布专员技能"""
        return cls.SEO_SKILLS + cls.PUBLISHING_SKILLS

# ========== 3. 定义状态 ==========

class ContentState(TypedDict):
    messages: Annotated[list, add_messages]
    topic: str
    research_data: dict
    outline: dict
    draft_content: str
    seo_optimized: str
    metadata: dict
    published: bool

# ========== 4. 创建专业化代理 ==========

def research_agent(state: ContentState):
    """研究代理：使用研究技能"""
    llm = ChatOpenAI(model="gpt-4o")
    skills = ContentSkillRegistry.get_researcher_skills()
    llm_with_tools = llm.bind_tools(skills)

    prompt = f"研究主题：{state['topic']}\n执行深度研究和竞品分析"
    response = llm_with_tools.invoke(prompt)

    return {
        "messages": [response],
        "research_data": {
            "sources": ["来源..."],
            "key_findings": ["发现..."]
        }
    }

def writer_agent(state: ContentState):
    """写作代理：使用写作技能"""
    llm = ChatAnthropic(model="claude-sonnet-4-5")
    skills = ContentSkillRegistry.get_writer_skills()
    llm_with_tools = llm.bind_tools(skills)

    research = state["research_data"]
    prompt = f"""基于研究数据创作文章：
    主题：{state['topic']}
    研究发现：{research}

    步骤：
    1. 使用 generate_outline 生成大纲
    2. 撰写完整文章
    3. 使用 check_grammar 检查语法
    """

    response = llm_with_tools.invoke(prompt)

    return {
        "messages": [response],
        "draft_content": "文章草稿内容..."
    }

def seo_agent(state: ContentState):
    """SEO 代理：使用 SEO 技能"""
    llm = ChatOpenAI(model="gpt-4o")
    skills = ContentSkillRegistry.get_seo_specialist_skills()
    llm_with_tools = llm.bind_tools(skills)

    draft = state["draft_content"]
    prompt = f"""优化 SEO：
    内容：{draft}

    步骤：
    1. 使用 analyze_seo 分析当前 SEO
    2. 优化内容和关键词密度
    3. 使用 generate_meta_tags 生成元标签
    """

    response = llm_with_tools.invoke(prompt)

    return {
        "messages": [response],
        "seo_optimized": "SEO 优化后的内容...",
        "metadata": {
            "title": "优化的标题",
            "description": "描述",
            "keywords": ["关键词1", "关键词2"]
        }
    }

def publisher_agent(state: ContentState):
    """发布代理：使用发布技能"""
    llm = ChatOpenAI(model="gpt-4o")
    skills = ContentSkillRegistry.get_publisher_skills()
    llm_with_tools = llm.bind_tools(skills)

    content = state["seo_optimized"]
    metadata = state["metadata"]

    prompt = f"""发布内容：
    内容：{content}
    元数据：{metadata}

    步骤：
    1. 使用 publish_to_cms 发布到网站
    2. 使用 schedule_social_media 安排社交媒体推广
    """

    response = llm_with_tools.invoke(prompt)

    return {
        "messages": [response],
        "published": True
    }

# ========== 5. 构建工作流 ==========

workflow = StateGraph(ContentState)

# 添加节点（每个代理使用不同的技能组合）
workflow.add_node("research", research_agent)
workflow.add_node("writer", writer_agent)
workflow.add_node("seo", seo_agent)
workflow.add_node("publisher", publisher_agent)

# 定义流程
workflow.add_edge(START, "research")
workflow.add_edge("research", "writer")
workflow.add_edge("writer", "seo")
workflow.add_edge("seo", "publisher")
workflow.add_edge("publisher", END)

# 编译
app = workflow.compile()

# ========== 6. 使用内容创作平台 ==========

result = app.invoke({
    "messages": [],
    "topic": "人工智能在内容创作中的应用"
})

print("内容创作流程完成！")
print(f"主题：{result['topic']}")
print(f"是否已发布：{result.get('published', False)}")
print(f"元数据：{result.get('metadata', {})}")

# 执行流程：
# 1. 研究代理：使用 web_research + competitor_analysis
# 2. 写作代理：使用 generate_outline + check_grammar
# 3. SEO 代理：使用 analyze_seo + generate_meta_tags
# 4. 发布代理：使用 publish_to_cms + schedule_social_media

# 优势：
# - 每个代理拥有精准的技能组合
# - 技能可复用（如 check_grammar 在多个代理间共享）
# - 易于维护（更新技能定义即可）
# - 易于扩展（添加新技能或新代理）
# - 专业化分工（研究、写作、SEO、发布）

"""
技能版本管理 - 支持多版本和 A/B 测试
"""
from typing import Dict, Callable
from enum import Enum

class SkillVersion(Enum):
    V1 = "v1"
    V2 = "v2"
    BETA = "beta"

class VersionedSkillRegistry:
    """支持版本管理的技能注册表"""

    def __init__(self):
        self._skills: Dict[str, Dict[SkillVersion, Callable]] = {}
        self._default_versions: Dict[str, SkillVersion] = {}

    def register(self, name: str, skill: Callable, version: SkillVersion):
        """注册技能的特定版本"""
        if name not in self._skills:
            self._skills[name] = {}
        self._skills[name][version] = skill

    def set_default_version(self, name: str, version: SkillVersion):
        """设置技能的默认版本"""
        self._default_versions[name] = version

    def get(self, name: str, version: SkillVersion = None):
        """获取技能（可指定版本）"""
        if name not in self._skills:
            return None

        # 使用指定版本或默认版本
        target_version = version or self._default_versions.get(name, SkillVersion.V1)
        return self._skills[name].get(target_version)

# 创建版本化注册表
versioned_registry = VersionedSkillRegistry()

# 注册不同版本的技能
@tool
def search_v1(query: str) -> dict:
    """搜索技能 V1：基础实现"""
    return {"results": ["结果1", "结果2"]}

@tool
def search_v2(query: str, filters: dict = None) -> dict:
    """搜索技能 V2：增强版，支持过滤"""
    # 新功能：支持过滤器
    return {
        "results": ["精准结果1", "精准结果2"],
        "filters_applied": filters or {}
    }

# 注册两个版本
versioned_registry.register("search", search_v1, SkillVersion.V1)
versioned_registry.register("search", search_v2, SkillVersion.V2)

# 设置默认版本
versioned_registry.set_default_version("search", SkillVersion.V1)

# A/B 测试：部分代理使用新版本
def create_ab_test_agent(agent_id: str, use_beta: bool = False):
    """创建 A/B 测试代理"""
    version = SkillVersion.V2 if use_beta else SkillVersion.V1
    search_skill = versioned_registry.get("search", version)

    def agent(state):
        llm = ChatOpenAI(model="gpt-4o")
        llm_with_tools = llm.bind_tools([search_skill])
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

    return agent

# 使用：50% 流量使用 V2
agent_a = create_ab_test_agent("agent_a", use_beta=False)  # V1
agent_b = create_ab_test_agent("agent_b", use_beta=True)   # V2

# 优势：
# - 渐进式更新（不影响现有代理）
# - A/B 测试（评估新版本效果）
# - 回滚能力（问题时切换回旧版本）
# - 灵活迁移（逐步迁移到新版本）

# 使用配置文件管理技能分配
skill_config = {
    "roles": {
        "researcher": {
            "skills": ["search_internet", "search_database", "analyze_data"],
            "model": "gpt-4o"
        },
        "writer": {
            "skills": ["generate_outline", "check_grammar", "format_report"],
            "model": "claude-sonnet-4-5"
        },
        "seo_specialist": {
            "skills": ["analyze_seo", "generate_meta_tags"],
            "model": "gpt-4o"
        }
    }
}

def load_agent_skills(role: str):
    """从配置加载代理技能"""
    config = skill_config["roles"].get(role, {})
    skill_names = config.get("skills", [])
    return skill_registry.get_skills(skill_names)

# 优势：
# - 集中配置，易于管理
# - 可外部化（YAML、JSON）
# - 便于动态调整
# - 支持环境差异（dev vs prod）

# 技能单元测试示例
import pytest

def test_search_internet():
    """测试搜索技能"""
    result = search_internet.invoke({"query": "LangChain"})
    assert "results" in result
    assert len(result["results"]) > 0

def test_analyze_data():
    """测试数据分析技能"""
    result = analyze_data.invoke({
        "data": "sample data",
        "analysis_type": "summary"
    })
    assert "insights" in result
    assert "score" in result
    assert 0 <= result["score"] <= 1

# 集成测试：测试技能组合
def test_skill_combination():
    """测试技能组合"""
    # 搜索 → 分析 → 格式化
    search_result = search_internet.invoke({"query": "test"})
    analysis = analyze_data.invoke({"data": str(search_result)})
    report = format_report.invoke({"data": analysis})
    assert report is not None

# 权限装饰器
from functools import wraps

def require_permission(permission: str):
    """技能权限装饰器"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 检查权限
            agent_permissions = kwargs.get("_permissions", [])
            if permission not in agent_permissions:
                raise PermissionError(f"需要权限: {permission}")
            return func(*args, **kwargs)
        wrapper.required_permission = permission
        return wrapper
    return decorator

@tool
@require_permission("admin")
def delete_data(data_id: str) -> bool:
    """删除数据（需要管理员权限）"""
    # 执行删除
    return True

# 为代理分配权限
researcher_permissions = ["search", "read"]
admin_permissions = ["search", "read", "write", "admin"]

# 运行时动态注册技能
class DynamicSkillRegistry(SkillRegistry):
    """支持动态添加技能"""

    def add_skill_at_runtime(self, skill: BaseTool, category: str):
        """运行时添加技能"""
        self.register(skill, category)
        print(f"已添加技能: {skill.name} ({category})")

    def remove_skill(self, name: str):
        """移除技能"""
        if name in self._skills:
            del self._skills[name]
            print(f"已移除技能: {name}")

# 使用
registry = DynamicSkillRegistry()

# 动态添加新技能
@tool
def new_feature(param: str) -> str:
    """新功能技能"""
    return f"处理: {param}"

registry.add_skill_at_runtime(new_feature, "experimental")