LangGraph Send API

langgraph_send.py

import operator
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from typing import Annotated, TypedDict
import openai


class State(TypedDict):
    topics: list[str]
    queries: Annotated[list[str], operator.add]


class CreateQueryState(TypedDict):
    topic: str


def initial_node(state: State) -> State:
    # Do some initialization here
    return {}


async def create_query(state: CreateQueryState) -> State:
    # Do some LLM call here
    topic = state["topic"]
    response = await openai.ChatCompletion.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Create a web search query for this topic"},
            {"role": "user", "content": f"Topic: {state['topic']}"},
        ],
    )
    return {"queries": [response.choices[0].message.content]}


async def synthesize_queries(state: State) -> State:
    # Do some LLM call here
    queries = state["queries"]
    response = await openai.ChatCompletion.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": "Synthesize the queries into a single meaningful query",
            },
            {"role": "user", "content": f"Queries: {queries}"},
        ],
    )
    return {"final_query": response.choices[0].message.content}


def fan_out_to_nodes(state: State) -> State:
    return [Send("create_query", {"topic": topic}) for topic in state["topics"]]


graph = StateGraph(dict)
graph.add_node("initial_node", initial_node)
graph.add_node("create_query", create_query)
graph.add_conditional_edges(
    "initial_node",
    fan_out_to_nodes,
    ["create_query"],
)
graph.add_edge("create_query", "synthesize_queries")
graph.add_edge("synthesize_queries", END)