Branch node

from pydantic import BaseModel, Field

from ibm_watsonx_orchestrate.flow_builder.flows import (
    Flow, flow, START, END, Branch
)

class Pet(BaseModel):
    kind: str = Field(description="the kind of pet: dog or cat")

class PetFacts(BaseModel):
    facts: list[str] = Field(description="A list of facts about the pet")

@flow(
        name = "get_pet_facts",
        input_schema = Pet,
        output_schema = PetFacts
)
def build_get_pet_facts_flow(aflow: Flow) -> Flow:
    """ Based on the request, we will return the list of facts about the pet. A pet can be either a cat or a dog. """
    
    dog_fact_node = aflow.tool("getDogFact")
    cat_fact_node = aflow.tool("getCatFact")

    # create a branch
    check_pet_kind_branch: Branch = aflow.branch(evaluator="flow.input.kind.strip().lower() == 'dog'")
    aflow.edge(START, check_pet_kind_branch)

    # edges are automatically added when cases are added
    check_pet_kind_branch.case(True, dog_fact_node).case(False, cat_fact_node)

    aflow.edge(dog_fact_node, END)
    aflow.edge(cat_fact_node, END)

    return aflow