第八章：跨框架对比——状态管理哲学

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph
import operator

# 定义状态的"白板"：多个类型化通道
class AgentState(TypedDict):
    # 消息历史：追加模式（reducer = operator.add）
    messages: Annotated[list, operator.add]
    
    # 当前任务：覆盖模式（默认 reducer = 最后一次写入覆盖）
    current_task: str
    
    # 搜索结果：追加模式
    search_results: Annotated[list, operator.add]
    
    # 进度计数器：取最大值（用于多节点并行时合并）
    progress: Annotated[int, max]
    
    # 最终答案：覆盖模式
    final_answer: str

import operator

# ===== 内建 Reducer 策略 =====

# 1. 追加（消息历史、结果收集）
messages: Annotated[list, operator.add]
# 写入 ["新消息"] → 旧值 ["旧消息1", "旧消息2"] + ["新消息"]
# 最终值 = ["旧消息1", "旧消息2", "新消息"]

# 2. 覆盖（当前状态、最新值）
current_step: str  # 不加 Annotated 默认覆盖
# 写入 "step_3" → 无论旧值是什么，直接变成 "step_3"

# 3. 取最大/最小值（进度、评分）
progress: Annotated[int, max]
# 节点A写入 5，节点B写入 3 → 合并后 = 5

# ===== 自定义 Reducer =====

# 4. 去重追加（搜索结果去重）
def dedup_add(existing: list, new: list) -> list:
    """追加新元素，但跳过已存在的"""
    seen = set(item["id"] for item in existing)
    return existing + [item for item in new if item["id"] not in seen]

results: Annotated[list, dedup_add]

# 5. 有界追加（只保留最近N条）
def bounded_add(existing: list, new: list, max_size: int = 100) -> list:
    """追加但保持总量不超过 max_size"""
    combined = existing + new
    return combined[-max_size:]  # 保留最近的

# 6. 投票合并（多节点投票决策）
def majority_vote(existing: dict, new: dict) -> dict:
    """合并多个节点的投票，返回当前票数"""
    merged = {**existing}
    for key, votes in new.items():
        merged[key] = merged.get(key, 0) + votes
    return merged

votes: Annotated[dict, majority_vote]

# 节点只读写自己关心的通道
def search_node(state: AgentState) -> dict:
    """搜索节点：读 current_task，写 search_results"""
    query = state["current_task"]          # 只读自己需要的
    results = do_search(query)
    return {"search_results": results}     # 只写自己负责的
    # 框架自动用对应 reducer 合并

def analyze_node(state: AgentState) -> dict:
    """分析节点：读 search_results，写 final_answer"""
    data = state["search_results"]         # 读搜索结果
    answer = analyze(data)
    return {"final_answer": answer}        # 覆盖写最终答案

# 图定义
graph = StateGraph(AgentState)
graph.add_node("search", search_node)
graph.add_node("analyze", analyze_node)
graph.add_edge("search", "analyze")

# LangGraph 内部用 BinaryOperatorAggregate 实现 Channel
class BinaryOperatorAggregate:
    """Channel 的内部实现：值 + 合并函数"""
    
    def __init__(self, value=None, operator=None):
        self.value = value
        self.operator = operator  # (旧值, 新值) -> 合并值
    
    def update(self, new_value):
        if self.value is None:
            self.value = new_value
        else:
            self.value = self.operator(self.value, new_value)
    
    def get(self):
        return self.value

# 三个搜索节点并行，结果通过 reducer 自动合并
graph.add_edge(START, "search_web")
graph.add_edge(START, "search_docs")
graph.add_edge(START, "search_db")
# 三个节点各自返回 {"findings": [...]}
# reducer (operator.add) 自动合并为完整列表

# IDE 能自动补全，写错 key 名立刻报错
def bad_node(state: AgentState) -> dict:
    x = state["nonexistent_key"]  # ← IDE 警告：key 不存在
    return {"typo_key": 123}      # ← IDE 警告：不在 State 定义中

# 简单对话 Agent 用 Channel 模式——显得很重
class SimpleState(TypedDict):
    messages: Annotated[list, operator.add]

graph = StateGraph(SimpleState)
graph.add_node("chat", chat_node)
graph.add_edge(START, "chat")
graph.add_edge("chat", END)
app = graph.compile()

# 对比 tRPC-Agent——三行搞定
agent := NewAgent(llm, tools)
agent.Run(userMessage)

from crewai import Agent, Task, Crew

# Task 的输出自动传递给后续 Task
research_task = Task(
    description="Research the topic: {topic}",
    expected_output="A comprehensive research report with key findings",
    agent=researcher
)

writing_task = Task(
    description="Write an article based on the research findings",
    # context 参数显式声明依赖哪些前置 Task 的输出
    context=[research_task],
    expected_output="A well-written article of 1000 words",
    agent=writer
)

review_task = Task(
    description="Review the article for accuracy and quality",
    context=[research_task, writing_task],  # 可以引用多个前置
    expected_output="Review feedback with specific suggestions",
    agent=editor
)

# CrewAI 内部的输出传递机制
class TaskOutput:
    raw: str              # 原始文本输出（总是有）
    pydantic: BaseModel   # 结构化输出（如果定义了 output_pydantic）
    json_dict: dict       # JSON 输出（如果定义了 output_json）
    agent: str            # 哪个 Agent 产出的
    
    # 下游 Task 的处理逻辑（简化）：
    # 1. 获取 context 中所有前置 Task 的 TaskOutput
    # 2. 把每个 TaskOutput.raw 拼接成上下文字符串
    # 3. 注入到下游 Task 的 prompt 中
    # 4. Agent 基于"任务描述 + 前置输出上下文"执行

# 如果第 5 个 Task 执行到一半崩了...
crew = Crew(tasks=[t1, t2, t3, t4, t5, t6])
result = crew.kickoff()  # t5 crash，前4个的结果全丢了

# 没有内建机制从 t5 恢复
# 必须自己在外部保存中间结果

# 想实现"写→审→改→再审→再改"的循环
# 纯 Task 链做不到——它是线性的，不能回头

# 只能用 CrewAI Flow 弥补（但那就不是纯 Task 传递模式了）

from crewai.flow.flow import Flow, start, listen, router

class ContentFlow(Flow):
    # Flow 有自己的 state（一个扁平 dict）
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    research_result: str = ""
    draft: str = ""
    revision_count: int = 0
    
    @start()
    def research(self):
        result = research_crew.kickoff()
        self.research_result = result.raw
        return result
    
    @listen(research)
    def write_draft(self):
        draft = writing_crew.kickoff(inputs={"research": self.research_result})
        self.draft = draft.raw
        return draft
    
    @router(write_draft)
    def quality_check(self):
        score = evaluate_quality(self.draft)
        if score > 0.8 or self.revision_count >= 3:
            return "publish"
        self.revision_count += 1
        return "revise"
    
    @listen("revise")
    def revise_draft(self):
        # 可以回到 write_draft 形成循环
        revised = revision_crew.kickoff(inputs={"draft": self.draft})
        self.draft = revised.raw
        return revised
    
    @listen("publish")
    def publish(self):
        return self.draft

from langgraph.graph import StateGraph
from typing import TypedDict, Annotated
import operator

class ResearchState(TypedDict):
    # 多个搜索节点并行，结果全部追加
    findings: Annotated[list, operator.add]
    # 多个评分节点并行，取最高分
    confidence: Annotated[float, max]
    # 多个节点投票
    votes: Annotated[dict, merge_votes]

graph = StateGraph(ResearchState)
graph.add_node("search_web", search_web_node)
graph.add_node("search_docs", search_docs_node)
graph.add_node("search_db", search_db_node)

# 三个节点从同一个起点出发，并行执行
graph.add_edge(START, "search_web")
graph.add_edge(START, "search_docs")
graph.add_edge(START, "search_db")

# 并行汇合点
graph.add_edge("search_web", "synthesize")
graph.add_edge("search_docs", "synthesize")
graph.add_edge("search_db", "synthesize")

# 执行时：
# search_web  返回 {"findings": [r1, r2], "confidence": 0.8}
# search_docs 返回 {"findings": [r3], "confidence": 0.6}
# search_db   返回 {"findings": [r4, r5], "confidence": 0.9}
#
# 框架自动合并：
# findings = [r1, r2] + [r3] + [r4, r5] = [r1, r2, r3, r4, r5]
# confidence = max(0.8, 0.6, 0.9) = 0.9

# Sequential：严格顺序，不可能并发
crew = Crew(
    agents=[researcher, writer, editor],
    tasks=[research_task, write_task, edit_task],
    process=Process.sequential
)

# Hierarchical：Manager 串行分配
# Manager 一次只把一个子任务交给一个 Agent
# 结果汇总到 Manager，Worker 之间不通信
crew = Crew(
    agents=[researcher, writer, editor],
    tasks=[complex_task],
    process=Process.hierarchical,
    manager_llm=ChatOpenAI(model="gpt-4")
)

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph

# 1. 配置 checkpoint 后端
checkpointer = SqliteSaver.from_conn_string("checkpoints.db")
# 也支持 PostgreSQL、Redis、自定义后端

# 2. 编译图时传入 checkpointer
app = graph.compile(checkpointer=checkpointer)

# 3. 执行——每个节点执行后自动保存快照
config = {"configurable": {"thread_id": "research-task-001"}}
result = app.invoke({"messages": [user_msg]}, config=config)

# 4. 恢复：O(1) 时间加载最新快照
state = app.get_state(config)
print(state.values)  # 所有 Channel 的当前值
print(state.next)    # 下一步应该执行哪个节点

# 5. 历史：获取所有 checkpoint
history = list(app.get_state_history(config))
# history[0] 是最新的，history[-1] 是最初的

# 6. Fork：从历史某点开始新分支
old_checkpoint = history[5]
fork_config = {"configurable": {
    "thread_id": "research-task-001-fork-a",
    "checkpoint_id": old_checkpoint.config["configurable"]["checkpoint_id"]
}}
# 从第5个checkpoint分叉，尝试不同策略
app.invoke({"messages": [alternative_instruction]}, fork_config)

# 7. Rollback：回退到某个历史状态
app.update_state(
    config,
    old_checkpoint.values,  # 用旧值覆盖当前值
    as_node=old_checkpoint.next[0]  # 指定从哪个节点恢复
)

from crewai.flow.flow import Flow, start, listen

class ResearchFlow(Flow):
    # Flow state（简单 dict）
    research_data: str = ""
    
    @start()
    def research(self):
        result = research_crew.kickoff()
        self.research_data = result.raw
        # @persist 会在这里保存 self 的所有属性
        return result

import json

class RobustFlow(Flow):
    checkpoint_path: str = "checkpoints/"
    
    @start()
    def step1(self):
        result = crew1.kickoff()
        # 手动保存
        self._save_checkpoint("step1", result.raw)
        return result
    
    def _save_checkpoint(self, step_name, data):
        with open(f"{self.checkpoint_path}/{step_name}.json", "w") as f:
            json.dump({"step": step_name, "data": data}, f)
    
    def _try_restore(self, step_name):
        path = f"{self.checkpoint_path}/{step_name}.json"
        if os.path.exists(path):
            with open(path) as f:
                return json.load(f)["data"]
        return None

class AgentState(TypedDict):
    # 覆盖型 Channel：永远只占一个值的空间
    current_task: str                    # O(1) 空间
    current_step: str                    # O(1) 空间
    final_answer: str                    # O(1) 空间
    
    # 有界追加：自定义 reducer 控制上限
    search_results: Annotated[list, bounded_add_10]  # 最多10条
    
    # 唯一可能增长的：messages
    messages: Annotated[list, operator.add]  # 需要额外管理

# Session 模式下，状态积累方式：
# [user_msg, llm_response, tool_call_1, tool_result_1, llm_response_2, 
#  tool_call_2, tool_result_2, ..., llm_final_answer]
# 每一条都是完整文本，token 线性增长

# Channel 模式下：
# messages: [user_msg, final_answer]  ← 只保留对话级消息
# search_results: [top_3_results]      ← 结构化数据，有界
# current_step: "synthesize"           ← 单值覆盖
# LLM 只需要读 messages + 当前节点关心的 Channel

# 10 个 Task 的流水线
# Task10 只看到 context 声明的前置 Task 输出
# 不是 Task1~Task9 的所有历史都堆起来

review_task = Task(
    description="Final review",
    context=[research_task, final_draft_task],  # 只看这两个的输出
    # 中间的搜索、整理、草稿修改过程——Task10 完全不知道
    expected_output="Final review notes"
)