使用路由构建多源知识库

"""
多源知识 router 示例

此示例演示多智能体系统的 router pattern。
router 对查询进行分类，将查询并行路由到专用智能体，
并将结果合成为组合响应。
"""

import operator
from typing import Annotated, Literal, TypedDict

from langchain.agents import create_agent
from langchain.chat_models import init_chat_model
from langchain.tools import tool
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from pydantic import BaseModel, Field


# 状态定义
class AgentInput(TypedDict):
    """每个 subagent 的简单输入状态。"""
    query: str


class AgentOutput(TypedDict):
    """每个 subagent 的输出。"""
    source: str
    result: str


class Classification(TypedDict):
    """单个路由决策：使用哪个查询调用哪个智能体。"""
    source: Literal["github", "notion", "slack"]
    query: str


class RouterState(TypedDict):
    query: str
    classifications: list[Classification]
    results: Annotated[list[AgentOutput], operator.add]
    final_answer: str


# 分类器的结构化输出 schema
class ClassificationResult(BaseModel):
    """将用户查询分类为智能体专用子问题的结果。"""
    classifications: list[Classification] = Field(
        description="要调用的智能体列表及其定向子问题"
    )


# 工具
@tool
def search_code(query: str, repo: str = "main") -> str:
    """在 GitHub 仓库中搜索代码。"""
    return f"在 {repo} 中找到与 '{query}' 匹配的代码：src/auth.py 中的身份验证中间件"


@tool
def search_issues(query: str) -> str:
    """搜索 GitHub issues 和 pull requests。"""
    return f"找到 3 个与 '{query}' 匹配的 issues：#142（API 身份验证文档）、#89（OAuth 流程）、#203（token 刷新）"


@tool
def search_prs(query: str) -> str:
    """搜索 pull requests 以查找实现细节。"""
    return f"PR #156 添加了 JWT 身份验证，PR #178 更新了 OAuth scopes"


@tool
def search_notion(query: str) -> str:
    """在 Notion workspace 中搜索文档。"""
    return f"找到文档：'API Authentication Guide'，涵盖 OAuth2 流程、API keys 和 JWT tokens"


@tool
def get_page(page_id: str) -> str:
    """按 ID 获取特定 Notion 页面。"""
    return f"页面内容：逐步身份验证设置说明"


@tool
def search_slack(query: str) -> str:
    """搜索 Slack 消息和线程。"""
    return f"在 #engineering 中找到讨论：'使用 Bearer tokens 进行 API 身份验证，刷新流程请参阅文档'"


@tool
def get_thread(thread_id: str) -> str:
    """获取特定 Slack 线程。"""
    return f"线程讨论了 API key 轮换的最佳实践"


# 模型和智能体
model = init_chat_model("openai:gpt-5.4")
router_llm = init_chat_model("openai:gpt-5.4-mini")

github_agent = create_agent(
    model,
    tools=[search_code, search_issues, search_prs],
    system_prompt=(
        "你是 GitHub 专家。通过搜索仓库、issues 和 pull requests，"
        "回答有关代码、API references 和实现细节的问题。"
    ),
)

notion_agent = create_agent(
    model,
    tools=[search_notion, get_page],
    system_prompt=(
        "你是 Notion 专家。通过搜索组织的 Notion workspace，"
        "回答有关内部流程、政策和团队文档的问题。"
    ),
)

slack_agent = create_agent(
    model,
    tools=[search_slack, get_thread],
    system_prompt=(
        "你是 Slack 专家。通过搜索团队成员分享知识和解决方案的"
        "相关线程和讨论来回答问题。"
    ),
)


# 工作流节点
def classify_query(state: RouterState) -> dict:
    """对查询进行分类，并确定要调用哪些智能体。"""
    structured_llm = router_llm.with_structured_output(ClassificationResult)

    result = structured_llm.invoke([
        {
            "role": "system",
            "content": """分析此查询，并确定要查询哪些知识库。
为每个相关来源生成针对该来源优化的定向子问题。

可用来源：
- github：代码、API references、实现细节、issues、pull requests
- notion：内部文档、流程、政策、团队 wiki
- slack：团队讨论、非正式知识分享、近期对话

仅返回与查询相关的来源。"""
        },
        {"role": "user", "content": state["query"]}
    ])

    return {"classifications": result.classifications}


def route_to_agents(state: RouterState) -> list[Send]:
    """根据分类结果 fan out 到智能体。"""
    return [
        Send(c["source"], {"query": c["query"]})
        for c in state["classifications"]
    ]


def query_github(state: AgentInput) -> dict:
    """查询 GitHub agent。"""
    result = github_agent.invoke({
        "messages": [{"role": "user", "content": state["query"]}]
    })
    return {"results": [{"source": "github", "result": result["messages"][-1].content}]}


def query_notion(state: AgentInput) -> dict:
    """查询 Notion agent。"""
    result = notion_agent.invoke({
        "messages": [{"role": "user", "content": state["query"]}]
    })
    return {"results": [{"source": "notion", "result": result["messages"][-1].content}]}


def query_slack(state: AgentInput) -> dict:
    """查询 Slack agent。"""
    result = slack_agent.invoke({
        "messages": [{"role": "user", "content": state["query"]}]
    })
    return {"results": [{"source": "slack", "result": result["messages"][-1].content}]}


def synthesize_results(state: RouterState) -> dict:
    """将所有智能体的结果组合成连贯答案。"""
    if not state["results"]:
        return {"final_answer": "未从任何知识来源找到结果。"}

    formatted = [
        f"**来自 {r['source'].title()}：**\n{r['result']}"
        for r in state["results"]
    ]

    synthesis_response = router_llm.invoke([
        {
            "role": "system",
            "content": f"""合成这些搜索结果，以回答原始问题："{state['query']}"

- 合并来自多个来源的信息，避免冗余
- 突出最相关且可操作的信息
- 标注不同来源之间的差异
- 保持响应简洁且结构清晰"""
        },
        {"role": "user", "content": "\n\n".join(formatted)}
    ])

    return {"final_answer": synthesis_response.content}


# 构建工作流
workflow = (
    StateGraph(RouterState)
    .add_node("classify", classify_query)
    .add_node("github", query_github)
    .add_node("notion", query_notion)
    .add_node("slack", query_slack)
    .add_node("synthesize", synthesize_results)
    .add_edge(START, "classify")
    .add_conditional_edges("classify", route_to_agents, ["github", "notion", "slack"])
    .add_edge("github", "synthesize")
    .add_edge("notion", "synthesize")
    .add_edge("slack", "synthesize")
    .add_edge("synthesize", END)
    .compile()
)


if __name__ == "__main__":
    result = workflow.invoke({
        "query": "如何对 API 请求进行身份验证？"
    })

    print("原始查询：", result["query"])
    print("\n分类：")
    for c in result["classifications"]:
        print(f"  {c['source']}: {c['query']}")
    print("\n" + "=" * 60 + "\n")
    print("最终答案：")
    print(result["final_answer"])