map_states.py

class GenerateAllPoems(MapStates):
    def states(
            self, state: State, context: ApplicationContext, inputs: Dict[str, Any]
    ) -> SyncOrAsyncGenerator[State]:
        for poem_type in state["poem_types"]:
            yield state.update(current_draft=None, poem_type=poem_type, feedback=[], num_drafts=0)

    def action(self, state: State, inputs: Dict[str, Any]) -> SubgraphType:
        graph = (
            GraphBuilder()
            .with_actions(
                edit,
                write,
                final_draft,
            )
            .with_transitions(
                ("write", "edit", Condition.expr(f"num_drafts < {state['max_drafts']}")),
                ("write", "final_draft"),
                ("edit", "final_draft", Condition.expr("len(feedback) == 0")),
                ("edit", "write"),
            )
        ).build()
        return RunnableGraph(graph=graph, entrypoint="write", halt_after=["final_draft"])

    def reduce(self, state: State, results: SyncOrAsyncGenerator[State]) -> State:
        proposals = []
        for output_state in results:
            proposals.append(output_state["final_draft"])
        return state.append(proposals=proposals)

    @property
    def writes(self) -> list[str]:
        return ["proposals"]

    @property
    def reads(self) -> list[str]:
        return ["poem_types", "poem_subject", "max_drafts"]