Parallelization

import os
import time
import operator
import uuid

from dotenv import load_dotenv
from typing import Annotated
from datetime import datetime
from pydantic import BaseModel
from typing_extensions import TypedDict
from langchain_openai import ChatOpenAI
from langchain.messages import HumanMessage
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.types import Send
from langgraph.func import entrypoint, task
from langgraph.checkpoint.memory import InMemorySaver
from IPython.display import Image, display
from concurrent.futures import ThreadPoolExecutor

# Load model configuration
_ = load_dotenv()

# Load model
model_name = "qwen3-max"
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url=os.getenv("DASHSCOPE_BASE_URL"),
    model=model_name,
    temperature=0.7,
)

# Create concurrent node a
def node_a(state: MessagesState):
    start_time = datetime.now()
    print(f"[node_a] Enter function time: {start_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}")

    # Use sleep to simulate time consumption
    time.sleep(2)

    end_time = datetime.now()
    print(f"[node_a] Exit function time: {end_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}")

    return {'messages': [HumanMessage(
        content=f'Node a ran for {round((end_time - start_time).total_seconds(), 3)} seconds'
    )]}

# Create concurrent node b
def node_b(state: MessagesState):
    start_time = datetime.now()
    print(f"[node_b] Enter function time: {start_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}")

    # Use sleep to simulate time consumption
    time.sleep(4)

    end_time = datetime.now()
    print(f"[node_b] Exit function time: {end_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}")

    return {'messages': [HumanMessage(
        content=f'Node b ran for {round((end_time - start_time).total_seconds(), 3)} seconds'
    )]}

# Create graph
builder = StateGraph(MessagesState)

# Add nodes
builder.add_node('node_a', node_a)
builder.add_node('node_b', node_b)

# Add edges
builder.add_edge(START, 'node_a')
builder.add_edge(START, 'node_b')
builder.add_edge('node_a', END)
builder.add_edge('node_b', END)

# Compile graph
my_graph = builder.compile(name='my-graph')
my_graph

# Invoke graph
response = my_graph.invoke({
    'messages': [HumanMessage(content='Execute node_a and node_b')]
})

for message in response['messages']:
    message.pretty_print()

# Concurrent execution unit
@task
def generate_paragraph(topic: str) -> str:
    response = llm.invoke([
        {"role": "system", "content": "You are an intelligent assistant who can write popular science articles"},
        {"role": "user", "content": (
            f"Write a paragraph about {topic}, approximately 100 words. "
            "What you write is just one paragraph of the entire article, "
            "so please focus on the given topic without introductory text")}
    ])
    return response.content

# Create persistent checkpoint
checkpointer = InMemorySaver()

# Workflow entry function, can concurrently call functions decorated by @task
@entrypoint(checkpointer=checkpointer)
def workflow(topics: list[str]) -> str:
    # Distribute topics to multiple task nodes for concurrent execution
    futures = [generate_paragraph(topic) for topic in topics]
    paragraphs = [f.result() for f in futures]
    return "\n\n".join(paragraphs)

# Run workflow
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
result = workflow.invoke(["What is sea urchin tofu", "How delicious is sea urchin tofu", "Ingredients and preparation of sea urchin tofu"], config=config)
print(result)

# Roles
class Roles(BaseModel):
    roles: list[str]

# Single role
class Role(BaseModel):
    role: str

# Single response
class Response(BaseModel):
    response: str

# Best response ID
class BestResponse(BaseModel):
    id: int

# Global context
class Overall(TypedDict):
    roles: list[str]
    responses: Annotated[list, operator.add]
    best_response: str

# Role prompt
role_prompt = "The goddess hasn't replied to your message again. As a {role}, how should you reply to the goddess in one sentence? Please return in JSON format with a response field"

# Best response prompt
best_response_prompt = """Below are reactions from several types of guys when the goddess doesn't reply to their messages.
Which reply do you think can best win back the goddess's heart? Please return the corresponding ID.
Note that the first reaction corresponds to ID 0. Return in JSON format with an id field.
Here are the guys' reactions:

{responses}"""

# [MAP] Use Send function to distribute roles
def continue_to_responses(state: Overall):
    return [ Send("generate_response", {"role": r}) for r in state["roles"] ]

# [MAP] Role response node: generate response for each role
def generate_response(state: Role):
    prompt = role_prompt.format(role=state["role"])
    response = llm.with_structured_output(Response).invoke(prompt)
    return {"responses": [response.response]}

# [REDUCE] Best response node: return the best response
def best_response(state: Overall):
    responses = "\n\n".join(state["responses"])
    prompt = best_response_prompt.format(responses=responses)
    response = llm.with_structured_output(BestResponse).invoke(prompt)
    return {"best_response": state["responses"][response.id]}

# Define Doge sub-graph output Schema
class DogeOutput(TypedDict):
    roles: list[str]
    responses: list[str]
    best_response: str

# Create graph
doge_builder = StateGraph(Overall, output_schema=DogeOutput)

# Add nodes
doge_builder.add_node("generate_response", generate_response)
doge_builder.add_node("best_response", best_response)

# Add edges
doge_builder.add_conditional_edges(START, continue_to_responses, ["generate_response"])
doge_builder.add_edge("generate_response", "best_response")
doge_builder.add_edge("best_response", END)

# Compile graph
doge_graph = doge_builder.compile(name='best-response')
doge_graph

# Invoke graph
roles = ["Dream Guy", "Troll", "Lovestruck Puppy", "Playboy", "Baby-faced Brother", "Socially Anxious Homebody",
         "Dominant CEO", "Tea-themed Guy 💅💅💅", "Artsy Long-haired Guy", "Kawaii Otaku"]
response = doge_graph.invoke({"roles": roles})

# Best response
print(f"Most liked reply by {model_name}: {response['best_response']}")

# All responses
for role, resp in zip(response["roles"], response["responses"]):
    print(f"【{role}】")
    print(resp)
    print()

# Define ReAct sub-graph global Schema
class ReActState(TypedDict):
    roles: list[str]
    n: int
    new_roles: list[str]
    responses: list[str]

# Divergently generate roles
def generate_roles(state: ReActState):
    roles = ", ".join(state["roles"])
    n = state["n"]
    prompt = (
        "Now we're playing a role-playing game about: how to win back the goddess when she doesn't reply to messages. "
        f"Reference role list: {roles}, create {n} different roles. "
        "You will play these roles later to reply to the goddess. "
        "Please return in JSON format with a roles field. "
        f"The roles field must be a pure string array with elements being the {n} newly created role names. "
        "Only return a JSON object without any explanation or other text."
    )
    res = llm.with_structured_output(Roles).invoke(prompt)
    return {"new_roles": res.roles}

def role_response(role: str) -> str:
    prompt = role_prompt.format(role=role)
    res = llm.with_structured_output(Response).invoke(prompt)
    return res.response

def divergent_thinking(state: ReActState):
    with ThreadPoolExecutor(max_workers=8) as executor:
        rs = list(executor.map(role_response, state["new_roles"]))
    return {"responses": rs}

# Define ReAct sub-graph output Schema
class ReActOutput(TypedDict):
    new_roles: list[str]
    responses: list[str]

# Create graph
react_builder = StateGraph(
    state_schema=ReActState,
    output_schema=ReActOutput
)

# Add nodes
react_builder.add_node('generate_roles', generate_roles)
react_builder.add_node('divergent_thinking', divergent_thinking)

# Add edges
react_builder.add_edge(START, 'generate_roles')
react_builder.add_edge('generate_roles', 'divergent_thinking')
react_builder.add_edge('divergent_thinking', END)

# Compile graph
react_graph = react_builder.compile(name='divergent-thinking')
react_graph

# Invoke graph
res = react_graph.invoke({"roles": roles, "n": 5})

assert len(res["new_roles"]) == len(res["responses"])

for role, resp in zip(res["new_roles"], res["responses"]):
    print(f"【{role}】")
    print(resp)
    print()

# Parent graph global Schema
class EntryGraphState(TypedDict):
    roles: list[str]
    n: int
    new_roles: list[str]
    responses: Annotated[list, operator.add]
    best_response: str

# Parent graph best response node: return best response and aggregate roles
def fa_best_response(state: EntryGraphState):
    responses = "\n\n".join(state["responses"])
    prompt = best_response_prompt.format(responses=responses)
    response = llm.with_structured_output(BestResponse).invoke(prompt)
    return {
        "roles": state["roles"] + state["new_roles"],
        "best_response": state["responses"][response.id]
    }

# Create parent graph
fa_builder = StateGraph(EntryGraphState)

# Add nodes
fa_builder.add_node("doge_graph", doge_builder.compile())
fa_builder.add_node("react_graph", react_builder.compile())
fa_builder.add_node("fa_best_response", fa_best_response)

# Add edges
fa_builder.add_edge(START, "doge_graph")
fa_builder.add_edge(START, "react_graph")
fa_builder.add_edge("doge_graph", "fa_best_response")
fa_builder.add_edge("react_graph", "fa_best_response")
fa_builder.add_edge("fa_best_response", END)

# Compile graph
fa_graph = fa_builder.compile(name='father-graph')

# Set xray parameter to 1 to display nested graph internal structure
display(Image(fa_graph.get_graph(xray=1).draw_mermaid_png()))

# Invoke graph
response = fa_graph.invoke({"roles": roles[3:6], "n": 5})

# Best response
print(f"Most liked reply by {model_name}: {response['best_response']}")

# All responses
assert len(response["roles"]) == len(response["responses"]), f'{len(response["roles"])} != {len(response["responses"])}'

for role, resp in zip(response["roles"], response["responses"]):
    print(f"【{role}】")
    print(resp)
    print()