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.base import Node
from fluksio.flow.nodes.wren import CANT, 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}


def test_a_reply_says_what_was_applied_and_the_rest_is_unhandled(fake_wren):
    node = WrenNode(
        requires=[MessageSpec(name="command", port="prompt", dtype=DType.STR)],
        provides=[
            MessageSpec(name="kitchen_light", dtype=DType.BOOL),
            MessageSpec(name="reply", dtype=DType.STR),
            MessageSpec(name="unhandled", dtype=DType.STR),
        ],
        params={"apply": True},
    )
    node.assign_flow("house", "control")
    fake_wren.answer = ([FakeCall("lights.on", {"id": "kitchen:light"})], 0.99)

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

    assert out == {"house.kitchen_light": True, "house.reply": "Kitchen light on."}

    # Nothing wired for the call, or nothing to call at all. Nothing reads
    # what is left, so the node says so itself.
    asked = "what is the weather tomorrow"
    for calls in ([FakeCall("lights.on", {"id": "garage:light"})], []):
        fake_wren.answer = (calls, 0.99)
        out = node.execute({"house.command": asked})
        assert out == {"house.unhandled": asked, "house.reply": CANT}

    # An assistant reading it answers instead.
    unhandled = MessageSpec(name="house.unhandled", port="unhandled", dtype=DType.STR)
    Pipeline([node, Node(lambda params, unhandled: None, requires=[unhandled])])
    assert node.execute({"house.command": asked}) == {"house.unhandled": asked}