app/backend/tests/flow/test_wren.py

"""The wren node: its wiring is the allowlist, and nothing is written unasked.

The model is an optional package, so ``fake_wren`` (conftest) answers in its
place: what the node does with an answer is what is under test here.
"""

from dataclasses import dataclass, field
from typing import Any

import pytest

from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes.wren import WrenNode, _value
from fluksio.flow.pipeline import Pipeline


@dataclass
class FakeCall:
    tool: str
    args: dict[str, Any] = field(default_factory=dict)
    delay_s: float | None = None


def _node(*targets: tuple[str, DType], **params: object) -> WrenNode:
    node = WrenNode(
        requires=[MessageSpec(name="command", port="prompt", dtype=DType.STR)],
        provides=[
            MessageSpec(name="actions", dtype=DType.LIST),
            *(MessageSpec(name=port, dtype=dtype) for port, dtype in targets),
        ],
        params=params,
    )
    node.assign_flow("house", "control")
    return node


def test_targets_of_one_area_and_kind_are_one_device(fake_wren):
    node = _node(
        ("living_room_setpoint", DType.FLOAT),
        ("living_room_boost", DType.BOOL),
        ("kitchen_lamp", DType.BOOL),
    )
    assert node.devices == {
        ("living room", "heat"): ["living_room_setpoint", "living_room_boost"],
        ("kitchen", "light"): ["kitchen_lamp"],
    }


@pytest.mark.parametrize("port", ["hallway_light", "kitchen_fridge", "light"])
def test_a_port_naming_no_trained_device_is_refused(fake_wren, port):
    with pytest.raises(ValueError, match="living_room") as refused:
        _node((port, DType.BOOL))
    assert port in str(refused.value)


def test_a_prompt_is_required(fake_wren):
    with pytest.raises(ValueError, match="prompt"):
        WrenNode(provides=[MessageSpec(name="kitchen_light", dtype=DType.BOOL)])


def test_only_the_prompt_wakes_it(fake_wren):
    node = WrenNode(
        requires=[
            MessageSpec(name="command", port="prompt", dtype=DType.STR),
            MessageSpec(name="outside", port="outdoor", dtype=DType.FLOAT),
        ],
    )
    assert {spec.port: spec.trigger for spec in node.input_ports} == {
        "prompt": True,
        "outdoor": False,
    }


@pytest.mark.parametrize(
    ("tool", "args", "dtype", "value"),
    [
        ("lights.on", {}, DType.BOOL, True),
        ("lights.on", {}, DType.FLOAT, 100.0),
        ("lights.off", {}, DType.STR, "off"),
        ("heating.set", {"temp": 22}, DType.FLOAT, 22.0),
        ("heating.set", {"temp": 22}, DType.BOOL, None),
        ("heating.boost", {}, DType.BOOL, True),
        ("heating.boost", {}, DType.INT, None),
        ("sensors.temperature", {"area": "kitchen"}, DType.FLOAT, None),
    ],
)
def test_a_call_means_what_its_target_type_can_hold(tool, args, dtype, value):
    assert _value(FakeCall(tool, args), dtype) == value


def test_applied_calls_reach_their_targets(fake_wren):
    node = _node(
        ("living_room_setpoint", DType.FLOAT),
        ("living_room_boost", DType.BOOL),
        apply=True,
    )
    Pipeline([node]).state["house.living_room_setpoint"] = 21.0
    fake_wren.answer = (
        [
            FakeCall("heating.set", {"id": "living room:heat", "temp": 23}),
            # The living room's light is filled in for the model, not wired.
            FakeCall("lights.on", {"id": "living room:light"}),
        ],
        0.98,
    )

    out = node.execute({"house.command": "set the heating to 23 degrees"})

    assert out is not None
    assert out["house.living_room_setpoint"] == 23.0
    assert "house.living_room_boost" not in out
    assert [(a["message"], a["applied"]) for a in out["house.actions"]] == [
        ("house.living_room_setpoint", True),
        ("", False),
    ]
    living = fake_wren.houses[0].areas["living room"]
    assert [(d.kind, d.state) for d in living] == [
        ("light", "off"),
        ("media", "off"),
        ("heat", "21C"),
    ]


@pytest.mark.parametrize(
    ("params", "confidence"),
    [({}, 0.99), ({"apply": True}, 0.5), ({"apply": True, "min_confidence": 1}, 0.99)],
)
def test_nothing_is_written_unless_applied_and_sure(fake_wren, params, confidence):
    node = _node(("kitchen_light", DType.BOOL), **params)
    fake_wren.answer = ([FakeCall("lights.on", {"id": "kitchen:light"})], confidence)

    out = node.execute({"house.command": "turn on the lights in the kitchen"})

    assert out == {
        "house.actions": [
            {
                "message": "house.kitchen_light",
                "value": True,
                "applied": False,
                "delay_s": None,
                "call": "lights.on()",
                "confidence": confidence,
            }
        ]
    }


def test_the_model_runs_on_one_blas_thread(fake_wren, monkeypatch):
    """About twice as fast for a model this small; numpy is loaded by `mlp`."""
    from threadpoolctl import threadpool_info

    threads: list[int] = []
    act = fake_wren.act

    def counted(command: str, house: Any) -> tuple[list[Any], float]:
        threads.extend(
            pool["num_threads"]
            for pool in threadpool_info()
            if pool["user_api"] == "blas"
        )
        return act(command, house)

    monkeypatch.setattr(fake_wren, "act", counted)
    _node(("kitchen_lamp", DType.BOOL)).execute({"house.command": "lights on"})

    assert threads
    assert set(threads) == {1}