"""
Based on the request, we will return the list of facts about the pet.
A pet can be either a cat or a dog.
"""
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_if_else",
input_schema = Pet,
output_schema = PetFacts
)
def build_get_pet_facts_if_else_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 with conditions
check_pet_kind_conditions_branch: Branch = aflow.conditions()
check_pet_kind_conditions_branch.condition(
expression="flow.input.kind.strip().lower() == 'dog'", to_node=dog_fact_node
).condition(expression="flow.input.kind.strip().lower() == 'cat'", to_node=cat_fact_node
).condition(to_node=dog_fact_node, default=True)
aflow.edge(START, check_pet_kind_conditions_branch)
aflow.edge(dog_fact_node, END)
aflow.edge(cat_fact_node, END)
return aflow