Timer node

from pydantic import BaseModel, Field
from ibm_watsonx_orchestrate.flow_builder.flows import END, Flow, flow, START
from .get_facts_about_numbers import get_facts_about_numbers
from .get_request_status import get_request_status

class InputtedNumber(BaseModel):
    number: int = Field(description="Inputted number from user")

class Attempt(BaseModel):
    atmp: int = Field(description="Represent each attemp to check if the request is finish", default=0)

class FlowOutput(BaseModel):
    info: str = Field(description="Fact about a number")

@flow(
    name="get_number_random_fact_flow",
    input_schema=InputtedNumber,
    output_schema=FlowOutput,
    description="A flow to get a random fact about a number"
)
def get_number_random_fact_flow(aflow: Flow) -> Flow:
    get_facts_about_numbers_node = aflow.tool(
        get_facts_about_numbers,
        input_schema=InputtedNumber,
        output_schema=FlowOutput
    )

    while_loop: Flow = aflow.loop(
        evaluator="not parent.get_request_status.input.attempt.atmp or parent.get_request_status.input.attempt.atmp < 5",
        input_schema=Attempt,
        output_schema=FlowOutput
    )

    timer_node = while_loop.timer(
        name="wait_1_sec",
        delay=1000,
        description="Wait for 1 second before polling again"
    )

    get_request_status_node = while_loop.tool(get_request_status)
    while_loop.sequence(START, get_request_status_node, timer_node, END)
    aflow.sequence(START, get_facts_about_numbers_node, while_loop, END)

    return aflow