import os
import getpass

from langchain_anthropic import ChatAnthropic

def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

_set_env("ANTHROPIC_API_KEY")

llm = ChatAnthropic(model="claude-3-5-sonnet-latest")

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from IPython.display import Image, display

# 圖狀態(tài)
class State(TypedDict):
    topic: str
    joke: str
    improved_joke: str
    final_joke: str

# 節(jié)點
def generate_joke(state: State):
    """第一次LLM調(diào)用，生成初始笑話"""
    msg = llm.invoke(f"Write a short joke about {state['topic']}")
    return {"joke": msg.content}

def check_punchline(state: State):
    """檢查笑話是否有笑點"""
    if"?"in state["joke"] or"!"in state["joke"]:
        return"Fail"
    return"Pass"

def improve_joke(state: State):
    """第二次LLM調(diào)用，改進笑話"""
    msg = llm.invoke(f"Make this joke funnier by adding wordplay: {state['joke']}")
    return {"improved_joke": msg.content}

def polish_joke(state: State):
    """第三次LLM調(diào)用，完善笑話"""
    msg = llm.invoke(f"Add a surprising twist to this joke: {state['improved_joke']}")
    return {"final_joke": msg.content}

# 構(gòu)建工作流
workflow = StateGraph(State)

# 添加節(jié)點
workflow.add_node("generate_joke", generate_joke)
workflow.add_node("improve_joke", improve_joke)
workflow.add_node("polish_joke", polish_joke)

# 添加邊連接節(jié)點
workflow.add_edge(START, "generate_joke")
workflow.add_conditional_edges(
    "generate_joke", check_punchline, {"Fail": "improve_joke", "Pass": END}
)
workflow.add_edge("improve_joke", "polish_joke")
workflow.add_edge("polish_joke", END)

# 編譯
chain = workflow.compile()

# 顯示工作流
display(Image(chain.get_graph().draw_mermaid_png()))

# 調(diào)用
state = chain.invoke({"topic": "cats"})
print("初始笑話：")
print(state["joke"])
print("\n--- --- ---\n")
if"improved_joke"in state:
    print("改進后的笑話：")
    print(state["improved_joke"])
    print("\n--- --- ---\n")
    print("最終笑話：")
    print(state["final_joke"])
else:
    print("笑話未通過質(zhì)量檢查——未檢測到笑點！")

# 圖狀態(tài)
class State(TypedDict):
    topic: str
    joke: str
    story: str
    poem: str
    combined_output: str

# 節(jié)點
def call_llm_1(state: State):
    """第一次LLM調(diào)用，生成笑話"""
    msg = llm.invoke(f"Write a joke about {state['topic']}")
    return {"joke": msg.content}

def call_llm_2(state: State):
    """第二次LLM調(diào)用，生成故事"""
    msg = llm.invoke(f"Write a story about {state['topic']}")
    return {"story": msg.content}

def call_llm_3(state: State):
    """第三次LLM調(diào)用，生成詩歌"""
    msg = llm.invoke(f"Write a poem about {state['topic']}")
    return {"poem": msg.content}

def aggregator(state: State):
    """將笑話、故事和詩歌合并為一個輸出"""
    combined = f"Here's a story, joke, and poem about {state['topic']}!\n\n"
    combined += f"STORY:\n{state['story']}\n\n"
    combined += f"JOKE:\n{state['joke']}\n\n"
    combined += f"POEM:\n{state['poem']}"
    return {"combined_output": combined}

# 構(gòu)建工作流
parallel_builder = StateGraph(State)

# 添加節(jié)點
parallel_builder.add_node("call_llm_1", call_llm_1)
parallel_builder.add_node("call_llm_2", call_llm_2)
parallel_builder.add_node("call_llm_3", call_llm_3)
parallel_builder.add_node("aggregator", aggregator)

# 添加邊連接節(jié)點
parallel_builder.add_edge(START, "call_llm_1")
parallel_builder.add_edge(START, "call_llm_2")
parallel_builder.add_edge(START, "call_llm_3")
parallel_builder.add_edge("call_llm_1", "aggregator")
parallel_builder.add_edge("call_llm_2", "aggregator")
parallel_builder.add_edge("call_llm_3", "aggregator")
parallel_builder.add_edge("aggregator", END)

# 編譯工作流
parallel_workflow = parallel_builder.compile()

# 顯示工作流
display(Image(parallel_workflow.get_graph().draw_mermaid_png()))

# 調(diào)用
state = parallel_workflow.invoke({"topic": "cats"})
print(state["combined_output"])

from typing_extensions import Literal
from langchain_core.messages import HumanMessage, SystemMessage

# 結(jié)構(gòu)化輸出的模式，用于路由邏輯
class Route(BaseModel):
    step: Literal["poem", "story", "joke"] = Field(
        None, descriptinotallow="路由過程中的下一步"
    )

# 為LLM添加結(jié)構(gòu)化輸出模式
router = llm.with_structured_output(Route)

# 狀態(tài)
class State(TypedDict):
    input: str
    decision: str
    output: str

# 節(jié)點
def llm_call_1(state: State):
    """寫一個故事"""
    result = llm.invoke(state["input"])
    return {"output": result.content}

def llm_call_2(state: State):
    """寫一個笑話"""
    result = llm.invoke(state["input"])
    return {"output": result.content}

def llm_call_3(state: State):
    """寫一首詩"""
    result = llm.invoke(state["input"])
    return {"output": result.content}

def llm_call_router(state: State):
    """根據(jù)輸入將任務(wù)路由到相應(yīng)的節(jié)點"""
    decision = router.invoke(
        [
            SystemMessage(
                cnotallow="根據(jù)用戶請求將輸入路由到故事、笑話或詩歌。"
            ),
            HumanMessage(cnotallow=state["input"]),
        ]
    )
    return {"decision": decision.step}

# 條件邊函數(shù)，根據(jù)路由決策將任務(wù)導(dǎo)向相應(yīng)節(jié)點
def route_decision(state: State):
    if state["decision"] == "story":
        return"llm_call_1"
    elif state["decision"] == "joke":
        return"llm_call_2"
    elif state["decision"] == "poem":
        return"llm_call_3"

# 構(gòu)建工作流
router_builder = StateGraph(State)

# 添加節(jié)點
router_builder.add_node("llm_call_1", llm_call_1)
router_builder.add_node("llm_call_2", llm_call_2)
router_builder.add_node("llm_call_3", llm_call_3)
router_builder.add_node("llm_call_router", llm_call_router)

# 添加邊連接節(jié)點
router_builder.add_edge(START, "llm_call_router")
router_builder.add_conditional_edges(
    "llm_call_router",
    route_decision,
    {
        "llm_call_1": "llm_call_1",
        "llm_call_2": "llm_call_2",
        "llm_call_3": "llm_call_3",
    },
)
router_builder.add_edge("llm_call_1", END)
router_builder.add_edge("llm_call_2", END)
router_builder.add_edge("llm_call_3", END)

# 編譯工作流
router_workflow = router_builder.compile()

# 顯示工作流
display(Image(router_workflow.get_graph().draw_mermaid_png()))

# 調(diào)用
state = router_workflow.invoke({"input": "Write me a joke about cats"})
print(state["output"])

from typing import Annotated, List
import operator

# 用于規(guī)劃的結(jié)構(gòu)化輸出模式
class Section(BaseModel):
    name: str = Field(descriptinotallow="報告部分的名稱")
    description: str = Field(descriptinotallow="本節(jié)涵蓋的主要主題和概念的簡要概述")

class Sections(BaseModel):
    sections: List[Section] = Field(descriptinotallow="報告的各個部分")

# 為LLM添加結(jié)構(gòu)化輸出模式
planner = llm.with_structured_output(Sections)

# 創(chuàng)建工作者節(jié)點
因為協(xié)調(diào)者-工作者工作流非常常見，LangGraph提供了Send API來支持這種模式。它允許你動態(tài)創(chuàng)建工作者節(jié)點，并為每個節(jié)點分配特定的輸入。每個工作者都有自己的狀態(tài)，并且所有工作者的輸出都會寫入一個共享的狀態(tài)鍵，協(xié)調(diào)者圖可以訪問這個鍵。這使得協(xié)調(diào)者能夠訪問所有工作者的輸出，并將它們整合成最終輸出。如下所示，我們遍歷一個部分列表，并將每個部分發(fā)送到一個工作者節(jié)點。

```python
from langgraph.constants import Send

# 圖狀態(tài)
class State(TypedDict):
    topic: str  # 報告主題
    sections: list[Section]  # 報告部分列表
    completed_sections: Annotated[list, operator.add]  # 所有工作者并行寫入此鍵
    final_report: str  # 最終報告

# 工作者狀態(tài)
class WorkerState(TypedDict):
    section: Section
    completed_sections: Annotated[list, operator.add]

# 節(jié)點
def orchestrator(state: State):
    """協(xié)調(diào)者生成報告計劃"""
    report_sections = planner.invoke(
        [
            SystemMessage(cnotallow="生成報告計劃"),
            HumanMessage(cnotallow=f"報告主題：{state['topic']}")
        ]
    )
    return {"sections": report_sections.sections}

def llm_call(state: WorkerState):
    """工作者撰寫報告的一部分"""
    section = llm.invoke(
        [
            SystemMessage(cnotallow="根據(jù)提供的名稱和描述撰寫報告部分。每個部分不加前言，使用Markdown格式。"),
            HumanMessage(cnotallow=f"部分名稱：{state['section'].name}，描述：{state['section'].description}")
        ]
    )
    return {"completed_sections": [section.content]}

def synthesizer(state: State):
    """將各部分整合成完整報告"""
    completed_report_sections = "\n\n---\n\n".join(state["completed_sections"])
    return {"final_report": completed_report_sections}

# 條件邊函數(shù)，為每個計劃部分分配工作者
def assign_workers(state: State):
    """為計劃中的每個部分分配工作者"""
    return [Send("llm_call", {"section": s}) for s in state["sections"]]

# 構(gòu)建工作流
orchestrator_worker_builder = StateGraph(State)

# 添加節(jié)點
orchestrator_worker_builder.add_node("orchestrator", orchestrator)
orchestrator_worker_builder.add_node("llm_call", llm_call)
orchestrator_worker_builder.add_node("synthesizer", synthesizer)

# 添加邊連接節(jié)點
orchestrator_worker_builder.add_edge(START, "orchestrator")
orchestrator_worker_builder.add_conditional_edges(
    "orchestrator", assign_workers, ["llm_call"]
)
orchestrator_worker_builder.add_edge("llm_call", "synthesizer")
orchestrator_worker_builder.add_edge("synthesizer", END)

# 編譯工作流
orchestrator_worker = orchestrator_worker_builder.compile()

# 顯示工作流
display(Image(orchestrator_worker.get_graph().draw_mermaid_png()))

# 調(diào)用
state = orchestrator_worker.invoke({"topic": "創(chuàng)建關(guān)于LLM擴展規(guī)律的報告"})

# 顯示最終報告
from IPython.display import Markdown
Markdown(state["final_report"])

# 圖狀態(tài)
class State(TypedDict):
    joke: str
    topic: str
    feedback: str
    funny_or_not: str

# 用于評估的結(jié)構(gòu)化輸出模式
class Feedback(BaseModel):
    grade: Literal["funny", "not funny"] = Field(descriptinotallow="判斷笑話是否有趣")
    feedback: str = Field(descriptinotallow="如果笑話不好笑，提供改進建議")

# 為LLM添加結(jié)構(gòu)化輸出模式
evaluator = llm.with_structured_output(Feedback)

# 節(jié)點
def llm_call_generator(state: State):
    """LLM生成笑話"""
    if state.get("feedback"):
        msg = llm.invoke(
            f"根據(jù)反饋生成關(guān)于{state['topic']}的笑話：{state['feedback']}"
        )
    else:
        msg = llm.invoke(f"生成關(guān)于{state['topic']}的笑話")
    return {"joke": msg.content}

def llm_call_evaluator(state: State):
    """LLM評估笑話"""
    grade = evaluator.invoke(f"評估笑話：{state['joke']}")
    return {"funny_or_not": grade.grade, "feedback": grade.feedback}

# 條件邊函數(shù)，根據(jù)評估器的反饋決定是否返回生成器或結(jié)束
def route_joke(state: State):
    """根據(jù)評估器的反饋決定是否返回生成器或結(jié)束"""
    if state["funny_or_not"] == "funny":
        return"Accepted"
    elif state["funny_or_not"] == "not funny":
        return"Rejected + Feedback"

# 構(gòu)建工作流
optimizer_builder = StateGraph(State)

# 添加節(jié)點
optimizer_builder.add_node("llm_call_generator", llm_call_generator)
optimizer_builder.add_node("llm_call_evaluator", llm_call_evaluator)

# 添加邊連接節(jié)點
optimizer_builder.add_edge(START, "llm_call_generator")
optimizer_builder.add_edge("llm_call_generator", "llm_call_evaluator")
optimizer_builder.add_conditional_edges(
    "llm_call_evaluator",
    route_joke,
    {
        "Accepted": END,
        "Rejected + Feedback": "llm_call_generator"
    }
)

# 編譯工作流
optimizer_workflow = optimizer_builder.compile()

# 顯示工作流
display(Image(optimizer_workflow.get_graph().draw_mermaid_png()))

# 調(diào)用
state = optimizer_workflow.invoke({"topic": "Cats"})
print(state["joke"])

from langchain_core.tools import tool

# 定義工具
@tool
def multiply(a: int, b: int) -> int:
    """乘法工具"""
    return a * b

@tool
def add(a: int, b: int) -> int:
    """加法工具"""
    return a + b

@tool
def divide(a: int, b: int) -> float:
    """除法工具"""
    return a / b

# 為LLM綁定工具
tools = [add, multiply, divide]
tools_by_name = {tool.name: tool for tool in tools}
llm_with_tools = llm.bind_tools(tools)