📜 Functional API

def task(
    __func_or_none__,
    *,
    name=None,
    retry_policy=None,     # 失败重试策略
    cache_policy=None,     # 结果缓存
    timeout=None,          # 超时控制
):
    """用 @task 装饰器定义一个 LangGraph 任务。

    重要规则：
    - task 只能在 entrypoint 或 StateGraph 内部调用
    - 调用 task 函数会【返回一个 future】（不是直接结果！）
    - future 可以并发收集 → 实现并行
    - 启用 checkpointer 时，输入输出必须可序列化
    """

class entrypoint:
    """用 @entrypoint 定义一个 LangGraph 工作流。

    被装饰函数必须接收【单个参数】（输入）。多参数用 dict 传。

    可注入的特殊参数（运行时自动注入）：
    - config   : RunnableConfig，运行时配置
    - previous : 同一 thread 上次调用的返回值（需 checkpointer）—— 跨调用记忆！
    - runtime  : Runtime 对象（含 context/store/writer）

    状态管理：
    - previous 让你能访问"上次运行的结果"
    - 若想让保存的值≠返回值，用 entrypoint.final
    """

# pip install langgraph
import time
from langgraph.func import task, entrypoint

# ============ ① @task 返回 future，实现并行 ============
@task
def slow_square(x: int) -> int:
    """一个慢任务（假装调 LLM）。"""
    time.sleep(1)   # 假装耗时
    return x * x

@entrypoint()
def parallel_squares(nums: list[int]) -> list[int]:
    """工作流：并行计算多个平方。"""
    # 关键：调用 task 返回的是 future，不会立即执行！
    futures = [slow_square(n) for n in nums]
    # 此时 3 个 task 已经【并发启动】，不是串行！

    # 收集 future 的结果（.result() 阻塞等待）
    results = [f.result() for f in futures]
    return results

start = time.time()
print(parallel_squares.invoke([1, 2, 3]))  # → [1, 4, 9]
print(f"耗时: {time.time()-start:.1f}s")    # → ~1s（并行），而非 3s（串行）


# ============ ② previous：跨调用的记忆 ============
from langgraph.checkpoint.memory import InMemorySaver

@entrypoint(checkpointer=InMemorySaver())
def counter(data, previous) -> int:
    """每次调用，在上次结果基础上累加。"""
    # previous = 上次同一 thread 的返回值（首次为 None）
    prev = previous or 0
    return prev + data["add"]

config = {"configurable": {"thread_id": "t1"}}
print(counter.invoke({"add": 5}, config))   # → 5   (previous=None)
print(counter.invoke({"add": 3}, config))   # → 8   (previous=5)
print(counter.invoke({"add": 10}, config))  # → 18  (previous=8)
# previous 实现了"跨调用记忆"，类似消息历史的效果


# ============ ③ 命令式的复杂控制流 ============
@task
def retrieve(query: str) -> str:
    return f"关于{query}的文档"

@task
def generate(query: str, context: str) -> str:
    return f"基于[{context}]回答{query}"

@entrypoint()
def rag(query: str) -> str:
    """命令式 RAG：比 StateGraph 更自然的写法。"""
    ctx = retrieve(query).result()   # 串行等待
    answer = generate(query, ctx).result()
    # 如果想判断"够不够"再检索，直接写 if：
    # if len(ctx) < 10: ctx = retrieve(query + " 更多").result()
    return answer

print(rag.invoke("什么是Agent"))  # → 基于文档回答

@entrypoint(checkpointer=InMemorySaver())
def workflow(data, previous):
    # previous 拿到的是【保存的】完整状态
    state = previous or {"count": 0, "history": []}
    state["count"] += 1
    state["history"].append(data)

    # 返回精简信息给调用者，但保存完整 state
    return entrypoint.final(
        value=f"已处理 {state['count']} 次",   # ← 返回给调用者
        save=state,                            # ← 保存到 checkpoint（下次 previous 拿这个）
    )