Send

from typing_extensions import TypedDict, Annotated
from operator import add
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
import random

class State(TypedDict):
  list: Annotated[list[str], add]

def random_node(state: State):
  # 随机生成边数量
  value = random.choice([1, 2, 3])
  # 自定义传入的下游 state，并且每条边传入的 state 数据不同
  return [Send("node", { "input": s }) for s in "abc"[:value]]

class SendInput(TypedDict):
  input: str

def node(state: SendInput):
  return {"list": [state["input"]]}

graph = StateGraph(State).add_node(node).add_conditional_edges(START, random_node).add_edge("node", END).compile()

result = graph.invoke({ "list": [] })
print(result)
# {'list': ['a', 'b', 'c']}
# {'list': ['a', 'b']}

import operator
from typing_extensions import Annotated, TypedDict
from my_openai.gpt_4_1_mini import get_gpt_4_1_mini
from langchain_core.prompts import PromptTemplate
from langgraph.types import Send
from langgraph.graph import END, StateGraph, START
from pydantic import BaseModel, Field

llm = get_gpt_4_1_mini()

class Subjects(BaseModel):
  '''笑话主题列表'''
  subjects: list[str] = Field(description="列表中的每个主题都是字符串")

class Joke(BaseModel):
  '''笑话'''
  joke: str = Field(description="笑话内容")

class BestJoke(BaseModel):
  '''最佳笑话下标'''
  index: int = Field(description="笑话的下标，从0开始")

class State(TypedDict):
  # 用户输入的主题
  topic: str
  # 笑话主题列表
  subjects: list[str]
  # 笑话列表
  jokes: Annotated[list, operator.add]
  # 最佳笑话
  best_selected_joke: str

class JokeState(TypedDict):
  subject: str

builder = StateGraph(State)

# 根据用户输入的主题，生成笑话主题列表
def generate_subjects(state: State):
  prompt = PromptTemplate.from_template(template="根据以下关键词: {topic}，生成一个使用逗号分隔的2至5个笑话主题列表").format(topic=state["topic"])
  response = llm.with_structured_output(Subjects).invoke(prompt)
  return { "subjects": response.subjects }

# 根据笑话主题，生成笑话
def generate_joke(state: JokeState):
  prompt = PromptTemplate.from_template(template="根据以下主题: {subject}，生成一个笑话").format(subject=state["subject"])
  response = llm.with_structured_output(Joke).invoke(prompt)
  return { "jokes": [response.joke] }

# 根据笑话主题列表，生成笑话
def continue_to_jokes(state: State):
  return [Send("generate_joke", { "subject": s }) for s in state["subjects"]]

# 根据用户输入的主题，和上面生成的笑话列表，选择最好笑的笑话
def best_joke(state: State):
  jokes = "\n".join(state["jokes"])
  prompt = PromptTemplate.from_template(template="以下是一些关于{topic}的笑话，选择最好笑的一个笑话，返回它的下标。\n\n笑话列表：\n{jokes}").format(topic=state["topic"], jokes=jokes)
  response = llm.with_structured_output(BestJoke).invoke(prompt)
  return { "best_selected_joke": state["jokes"][response.index] }

builder.add_node(generate_subjects)
builder.add_node(generate_joke)
builder.add_node(continue_to_jokes)
builder.add_node(best_joke)
builder.add_edge(START, "generate_subjects")
builder.add_conditional_edges("generate_subjects", continue_to_jokes)
builder.add_edge("generate_joke", "best_joke")
builder.add_edge("best_joke", END)
graph = builder.compile()

for s in graph.stream({ "topic": "动物" }):
  print(s)