graph/workflow.py

"""LangGraph 状态机编排 —— 把 Week2 的串行 analyze() 升级为并行图。

流程:
    START → router(意图路由)
      ├─ doc_analysis → [event, sentiment, summary, topic] 并行扇出 → synthesis → END
      └─ qa           → qa(可注入 RAG context)→ END

- 4 个分析 Agent 并行跑(LangGraph 自动并发),结果经 reducer 累加到 analyses。
- synthesis 等 4 个都完成后汇总成结构化报告(确定性合并,不再调模型)。
"""
import operator
from typing import Annotated, Optional, TypedDict

from langgraph.graph import StateGraph, START, END

from agents import route
from agents.analysts import AGENTS, DOC_ANALYSTS, QA
from rag import HybridRetriever, KnowledgeBase

# 检索器单例:首次用时加载已持久化的知识库并建 BM25(小库,一次性);
# 知识库变更(上传新文档)后调 refresh_retriever() 重建。
_retriever = None


def get_retriever():
    global _retriever
    if _retriever is None:
        _retriever = HybridRetriever(KnowledgeBase())
    return _retriever


def refresh_retriever():
    global _retriever
    _retriever = None


class FinState(TypedDict, total=False):
    text: str                                  # 输入文本/问题
    context: Optional[str]                     # RAG 检索资料(qa 流用)
    router: dict                               # 路由决策详情
    intent: str                                # doc_analysis | qa
    analyses: Annotated[list, operator.add]    # 并行分析结果(reducer 累加)
    report: dict                               # synthesis 汇总产物
    answer: dict                               # qa 回答
    sources: list                              # qa 检索到的来源(doc_id+score)


# ---------------- 节点 ----------------
def router_node(state: FinState) -> dict:
    d = route(state["text"])
    return {"intent": d["intent"], "router": d}


def make_analyst_node(name: str):
    """为某个分析 Agent 生成节点;只往 analyses 追加自己的结果。"""
    agent = AGENTS[name]

    def node(state: FinState) -> dict:
        return {"analyses": [agent.run(state["text"])]}

    return node


def qa_node(state: FinState) -> dict:
    """问答流:未显式给 context 时,自动走 RAG 混合检索注入。"""
    ctx = state.get("context")
    sources = []
    if ctx is None:
        try:
            hits = get_retriever().retrieve(state["text"], k=4)
            sources = [{"doc_id": h["meta"].get("doc_id"), "score": h["score"]} for h in hits]
            ctx = "\n\n".join(f"[{h['meta'].get('doc_id')}] {h['text']}" for h in hits) or None
        except Exception:
            ctx = None   # 知识库不可用/为空时退化为无检索问答
    return {"answer": QA.run(state["text"], context=ctx), "sources": sources}


def synthesis_node(state: FinState) -> dict:
    """把 4 个分析结果合并成一份结构化报告。"""
    by = {a["agent"]: a for a in state.get("analyses", [])}

    def res(name):
        a = by.get(name)
        return a["result"] if a and a.get("ok") else None

    ev, se, su, to = res("event"), res("sentiment"), res("summary"), res("topic")
    report = {
        "events": (ev or {}).get("events", []),
        "sentiment": (se or {}).get("sentiment"),
        "sentiment_reason": (se or {}).get("reason"),
        "summary": (su or {}).get("summary"),
        "key_points": (su or {}).get("key_points", []),
        "industry": (to or {}).get("industry"),
        "topics": (to or {}).get("topics", []),
        "entities": (to or {}).get("entities", []),
        "failed": [a["agent"] for a in state.get("analyses", []) if not a.get("ok")],
    }
    return {"report": report}


# ---------------- 路由 ----------------
def decide(state: FinState):
    """条件边:doc_analysis 扇出到 4 个分析节点;qa 走问答节点。"""
    if state["intent"] == "doc_analysis":
        return [a.name for a in DOC_ANALYSTS]   # 返回列表 = 并行扇出
    return "qa"


# ---------------- 组图 ----------------
def build_graph():
    g = StateGraph(FinState)
    g.add_node("router", router_node)
    for a in DOC_ANALYSTS:
        g.add_node(a.name, make_analyst_node(a.name))
    g.add_node("synthesis", synthesis_node)
    g.add_node("qa", qa_node)

    g.add_edge(START, "router")
    g.add_conditional_edges(
        "router", decide,
        [a.name for a in DOC_ANALYSTS] + ["qa"],
    )
    for a in DOC_ANALYSTS:
        g.add_edge(a.name, "synthesis")   # 4 个分析都汇入 synthesis(自动 join)
    g.add_edge("synthesis", END)
    g.add_edge("qa", END)
    return g.compile()


GRAPH = build_graph()


def analyze(text: str, context: Optional[str] = None) -> dict:
    """对外入口:与 Week2 agents.analyze 同签名,内部走 LangGraph。"""
    return GRAPH.invoke({"text": text, "context": context, "analyses": []})