from typing import TypedDict
from langchain_deepseek import ChatDeepSeek
from langgraph.graph import END, START, StateGraph
from dotenv import load_dotenv
#0. 配置模型
load_dotenv()
llm = ChatDeepSeek(
    model="deepseek-chat",
)

# 1. 定义图状态
class State(TypedDict):
    topic: str
    joke: str
    improved_joke: str
    final_joke: str

# 2. 定义节点函数
def generate_joke(state: State):
    '''
    第一个大模型调用，根据主题生成初始笑话
    '''
    topic = state['topic']
    msg = llm.invoke(f'写一个关于{topic}的简短笑话')
    return {
        'joke': msg.content
    }
def check_punchline(state: State):
    '''
    模拟门控函数——笑话中是否包含?或!
    '''
    joke = state['joke']
    if '?' in joke or '？' in joke:
        return "Fail" # 未能通过门控检查, 需要继续增强
    return "Pass"
def improve_joke(state: State):
    '''
    第二个大模型调用，通过添加文字游戏改进笑话
    '''
    joke = state['joke']
    msg = llm.invoke(f'通过添加文字游戏使笑话更有趣，当前笑话是: {joke}')
    return {
        'improved_joke': msg.content
    }
def polish_joke(state: State):
    '''
    第三个大模型调用，最终润色笑话，添加令人惊讶的转折
    '''
    improved_joke = state['improved_joke']
    msg = llm.invoke(f'为这个笑话添加一个令人惊讶的转折: {improved_joke}')
    return {
        'final_joke': msg.content
    }

#2. 定义边和图
workflow = StateGraph(State)
workflow.add_node('generate_joke', generate_joke)
workflow.add_node('improve_joke', improve_joke)
workflow.add_node('polish_joke', polish_joke)
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()

#3. 测试运行
state = chain.invoke({'topic': '小猫'})
print('初始笑话:')
print(state['joke'])
if 'improved_joke' in state:
    print('改进后笑话:')
    print(state['improved_joke'])
    print('最终笑话:')
    print(state['final_joke'])

from typing import TypedDict, Literal
from langchain_deepseek import ChatDeepSeek
from langgraph.graph import StateGraph, START, END
from langchain.messages import HumanMessage, SystemMessage
from dotenv import load_dotenv
load_dotenv()
llm = ChatDeepSeek(
    model="deepseek-chat",
)

class State(TypedDict):
    input: str
    decision: str
    output: str

# 1. 定义节点
def generate_story(state: State):
    '''
    写故事节点
    '''
    print('进入写故事处理逻辑')
    result = llm.invoke(state['input'])
    return {
        'output': result.content
    }
def generate_joke(state: State):
    '''
    写笑话节点
    '''
    print('进入写笑话处理逻辑')
    result = llm.invoke(state['input'])
    return {
        'output': result.content
    }
def generate_poetry(state: State):
    '''
    写诗歌节点
    '''
    print('进入写诗歌处理逻辑')
    result = llm.invoke(state['input'])
    return {
        'output': result.content
    }
class Classification(TypedDict):
    response_format: Literal['story', 'joke', 'poetry']
def llm_call_router(state: State):
    '''
    使用结构化输出将输入路由到适当的节点
    '''
    structed_llm = llm.with_structured_output(Classification)
    input_content = state['input']
    response = structed_llm.invoke([
        SystemMessage(cnotallow='''
            你是一个分类路由，根据用户的输入进行分类，分类结果是story, joke, poetry三者中的一种
        '''),
        HumanMessage(cnotallow=input_content)
    ])
    return {
        'decision': response['response_format']
    }

# 定义条件边函数
def route_decision(state: State):
    if state['decision'] == 'story':
        return 'llm_story'
    elif state['decision'] == 'joke':
        return 'llm_joke'
    elif state['decision'] == 'poetry':
        return 'llm_poetry'