from typing import TypedDict, NotRequired

class WorkflowState(TypedDict):
    user_query: str  # 用户原始查询（必须字段）
    tool_result: NotRequired[str]  # 工具调用结果（非必须字段）
    final_answer: NotRequired[str]  # 最终回答
    progress: NotRequired[int]  # 任务进度百分比

def node_function(state: WorkflowState):
    # 只需返回要更新的字段
    return {
        "tool_result": "处理结果",
        "progress": 50  # 只更新进度字段
    }

# LLM调用节点示例
def llm_node(state: WorkflowState):
    prompt = f"请回答：{state['user_query']}"
    response = llm.invoke(prompt)
    return {"llm_answer": response.content}

# 工具调用节点示例
def search_node(state: WorkflowState):
    result = search_tool.invoke(state["user_query"])
    return {"tool_result": result}

# 数据处理节点示例
def format_node(state: WorkflowState):
    formatted = format_output(state["llm_answer"])
    return {"final_answer": formatted}

from langgraph.graph import StateGraph, START, END

# 定义工作流
workflow = StateGraph(WorkflowState)

# 添加节点
workflow.add_node("parse_query", parse_query)
workflow.add_node("call_tool", call_tool)
workflow.add_node("generate_answer", generate_answer)

# 定义执行顺序
workflow.set_entry_point("parse_query")  # 入口
workflow.add_edge("parse_query", "call_tool")  # 普通边
workflow.add_edge("call_tool", "generate_answer")

# 添加条件边
def should_continue(state: WorkflowState):
    if state.get("needs_refinement"):
        return "generate_answer"
    return END

workflow.add_conditional_edges(
    "generate_answer",
    should_continue,
    {
        "refine": "parse_query",  # 需要重写
        END: END  # 完成
    }
)

# 编译
graph = workflow.compile()

from langchain_openai import ChatOpenAI
from dotenv import load_dotenv

load_dotenv()

# 初始化模型（兼容任意OpenAI兼容API）
llm = ChatOpenAI(
    model="deepseek-chat",
    temperature=0.3,
    api_key=os.getenv("API_KEY"),
    base_url=os.getenv("BASE_URL")
)

# 简单调用
response = llm.invoke("用一句话解释大语言模型")
print(response.content)

from langchain_core.prompts import ChatPromptTemplate

# 定义模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业的{domain}助手"),
    ("human", "请回答以下问题：{question}")
])

# 动态填充
chain = prompt | llm
response = chain.invoke({
    "domain": "人工智能",
    "question": "什么是过拟合？"
})

from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field

class Answer(BaseModel):
    answer: str = Field(description="问题的答案")
    confidence: float = Field(description="答案的可信度，0-1之间")
    sources: list[str] = Field(description="参考来源列表")

parser = JsonOutputParser(pydantic_object=Answer)
chain = prompt | llm | parser

result = chain.invoke({"question": "什么是大模型？"})
# result: {"answer": "...", "confidence": 0.95, "sources": [...]}

from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.chat_history import InMemoryChatMessageHistory

# 简单内存实现
history = InMemoryChatMessageHistory()
history.add_user_message("我叫张三")
history.add_ai_message("你好张三，有什么可以帮你？")
history.add_user_message("我叫李四")

messages = history.messages
# [HumanMessage(content="我叫张三"),
#  AIMessage(content="你好张三，有什么可以帮你？"),
#  HumanMessage(content="我叫李四")]

from langchain_core.tools import tool

@tool
def search_paper(topic: str) -> str:
    """搜索学术论文"""
    # 实际项目中连接外部API
    return f"找到3篇关于{topic}的论文"

@tool
def get_citation(paper_title: str) -> str:
    """获取论文引用信息"""
    return f"{paper_title} - 作者: XXX, 发表: 2024"

# 绑定工具
tools = [search_paper, get_citation]
llm_with_tools = llm.bind_tools(tools)

class WerewolfState(TypedDict):
    game_phase: str  # night_phase1, night_phase2, day_speech, vote
    players: list[Player]
    dead_players: list[str]
    wolf_targets: NotRequired[list[str]]
    seer_check_result: NotRequired[str]
    witch_potions: NotRequired[dict]
    speeches: NotRequired[list[Speech]]
    votes: NotRequired[dict]

# ✅ 推荐：最小必要字段，清晰命名
class TaskState(TypedDict):
    user_query: str
    context: NotRequired[str]
    current_step: NotRequired[str]
    result: NotRequired[str]
    history: NotRequired[list[dict]]  # 完整的操作历史

# ❌ 不推荐：字段过多、职责不清
class BadState(TypedDict):
    everything: str  # 什么都能存，但什么都说不清
    temp_data: NotRequired[any]  # 滥用any类型
    misc: NotRequired[dict]  # 大杂烩字段

def safe_node(state: WorkflowState):
    try:
        # 主要逻辑
        result = risky_operation(state["data"])
        return {"result": result, "error": None}
    except ValidationError as e:
        return {"error": f"验证失败: {e}"}
    except NetworkError as e:
        return {"error": f"网络错误: {e}", "retry_count": state.get("retry_count", 0) + 1}

from langgraph.callbacks import get_callback_handler

# 添加回调支持监控
graph = workflow.compile(
    callbacks=[get_callback_handler()]
)

# 或使用LangSmith进行追踪
os.environ["LANGSMITH_TRACING"] = "true"
os.environ["LANGSMITH_API_KEY"] = "your-api-key"

from langgraph.checkpoint.sqlite import SqliteSaver

# 使用SQLite持久化状态
checkpointer = SqliteSaver.from_conn_string(":memory:")
graph = workflow.compile(checkpointer=checkpointer)

# 从检查点恢复
result = graph.invoke(input, config={"configurable": {"thread_id": "session-123"}})