stroblme
5e5f1836cb
threadpoolctl scopes the limit to wren.act, so other numpy calls in the engine keep their threads.
354 lines
13 KiB
Python
354 lines
13 KiB
Python
"""Wren: a typed command in, values for the devices wired to it out.
|
|
|
|
`Wren <https://github.com/stroblme/wren>`_ is a small language model trained on
|
|
one job: read a house — its areas, the devices in each and their states — and
|
|
answer a command with calls on those devices. This node is the bridge. Its
|
|
output ports *are* the devices, named ``<area>_<kind>``, so the wiring is the
|
|
allowlist: the model can reach nothing that is not wired to the node.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
from threadpoolctl import threadpool_limits
|
|
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.nodes.base import Node
|
|
|
|
|
|
def _wren() -> Any:
|
|
"""The model package, which is optional.
|
|
|
|
Imported when a node is built rather than at boot, so an instance without
|
|
it still offers every other type, and a node placed without it says so
|
|
on the canvas straight away.
|
|
"""
|
|
try:
|
|
import wren
|
|
except ImportError:
|
|
raise RuntimeError("the wren node needs the model: pip install wren") from None
|
|
return wren
|
|
|
|
|
|
#: The last word of a target port, and the kind of device it stands for.
|
|
KIND_WORDS = {
|
|
"light": "light",
|
|
"lamp": "light",
|
|
"brightness": "light",
|
|
"setpoint": "heat",
|
|
"heating": "heat",
|
|
"heat": "heat",
|
|
"boost": "heat",
|
|
"media": "media",
|
|
"music": "media",
|
|
"tv": "media",
|
|
"shutter": "shutter",
|
|
"blind": "shutter",
|
|
"cover": "shutter",
|
|
"water": "water",
|
|
"watering": "water",
|
|
"sprinkler": "water",
|
|
}
|
|
|
|
#: What the node reads, and the types each may be. Only the prompt wakes it.
|
|
INPUTS = {
|
|
"prompt": {DType.STR},
|
|
"presence": {DType.STR},
|
|
"outdoor": {DType.FLOAT, DType.INT},
|
|
"dark": {DType.BOOL},
|
|
}
|
|
|
|
#: The output the calls are reported on. Every other output is a device.
|
|
ACTIONS = "actions"
|
|
|
|
_TARGET_TYPES = {DType.BOOL, DType.INT, DType.FLOAT, DType.STR}
|
|
|
|
#: The tools that act on each kind.
|
|
_TOOLS = {
|
|
"light": "lights",
|
|
"media": "media",
|
|
"heat": "heating",
|
|
"shutter": "shutters",
|
|
"water": "watering",
|
|
}
|
|
|
|
#: What a call sets a target to, by the target's type: (bool, number, text).
|
|
#: None is a call that type has no value for, which is then not applied.
|
|
#: `heating.set` carries its own number.
|
|
_VALUES: dict[str, tuple[bool | None, float | None, str | None]] = {
|
|
"lights.on": (True, 100, "on"),
|
|
"lights.off": (False, 0, "off"),
|
|
"media.play": (True, None, "play"),
|
|
"media.stop": (False, None, "stop"),
|
|
"heating.boost": (True, None, "boost"),
|
|
"heating.off": (False, None, "off"),
|
|
"shutters.up": (True, 100, "up"),
|
|
"shutters.down": (False, 0, "down"),
|
|
"shutters.stop": (None, None, "stop"),
|
|
"watering.start": (True, None, "on"),
|
|
"watering.stop": (False, None, "off"),
|
|
}
|
|
|
|
#: What every area held in training: a light, and in the living room, which
|
|
#: always existed, media and heating besides. A house missing them reads as
|
|
#: unfamiliar and the confidence drops with it, so they are filled in; a call
|
|
#: on one that nothing is wired to is reported and not applied.
|
|
_TRAINED = {"living room": ("light", "media", "heat")}
|
|
|
|
|
|
def _device(port: str, areas: dict[str, str]) -> tuple[str, str] | None:
|
|
"""The (area, kind) a port name stands for, or None."""
|
|
for slug, area in areas.items():
|
|
word = port.removeprefix(f"{slug}_")
|
|
if word != port and word in KIND_WORDS:
|
|
return area, KIND_WORDS[word]
|
|
return None
|
|
|
|
|
|
def _on(value: Any) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, (int, float)):
|
|
return value > 0
|
|
return str(value).strip().lower() in {"on", "up", "open", "play", "true"}
|
|
|
|
|
|
def _state(kind: str, values: list[Any]) -> str:
|
|
"""A device's state as the model reads it, from what its targets hold."""
|
|
if kind == "heat":
|
|
temps = [
|
|
v for v in values if isinstance(v, (int, float)) and not isinstance(v, bool)
|
|
]
|
|
return f"{round(temps[0])}C" if temps else "20C"
|
|
on = bool(values) and _on(values[0])
|
|
if kind == "shutter":
|
|
return "up" if on else "down"
|
|
return "on" if on else "off"
|
|
|
|
|
|
def _value(call: Any, dtype: DType) -> Any:
|
|
"""What ``call`` sets a target of this type to, or None."""
|
|
if call.tool == "heating.set":
|
|
temp = call.args.get("temp")
|
|
flag, number, text = (
|
|
None,
|
|
temp if isinstance(temp, (int, float)) else None,
|
|
None,
|
|
)
|
|
else:
|
|
flag, number, text = _VALUES.get(call.tool, (None, None, None))
|
|
if dtype is DType.BOOL:
|
|
return flag
|
|
if dtype is DType.STR:
|
|
return text
|
|
if number is None:
|
|
return None
|
|
return int(number) if dtype is DType.INT else float(number)
|
|
|
|
|
|
def _text(call: Any) -> str:
|
|
"""The call as the chat shows it, less the device it names."""
|
|
args = ", ".join(f"{k}={v}" for k, v in call.args.items() if k != "id")
|
|
return f"{call.tool}({args})"
|
|
|
|
|
|
class WrenNode(Node):
|
|
"""Turn the command on ``prompt`` into values for the devices wired to it.
|
|
|
|
Inputs: ``prompt`` (the command, and the only one that wakes the node),
|
|
``presence`` (the area the user is in), ``outdoor`` (°C) and ``dark``
|
|
(otherwise read off the clock: before 07:00 or from 18:00).
|
|
|
|
Outputs: ``actions``, a list of records, one per call and target, and
|
|
any number of targets named ``<area>_<kind>``. Targets of one area and kind
|
|
are one device to the model; each takes what the call means for its type.
|
|
A delayed call is held by the engine's work queue, like the Trigger
|
|
node's second value.
|
|
"""
|
|
|
|
class Params(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
apply: bool = Field(
|
|
default=False,
|
|
description="Write to the targets. Off, the node only reports what it would do.",
|
|
)
|
|
min_confidence: float = Field(
|
|
default=0.93,
|
|
ge=0,
|
|
le=1,
|
|
description="Below this, nothing is written; the calls are still reported.",
|
|
)
|
|
|
|
__slots__ = ("cfg", "devices")
|
|
|
|
# A call is a command to a device; doing it twice is not doing it once.
|
|
idempotent = False
|
|
|
|
def __init__(
|
|
self,
|
|
requires: MessageSpec | Iterable[MessageSpec] = (),
|
|
provides: MessageSpec | Iterable[MessageSpec] = (),
|
|
params: dict[str, Any] | None = None,
|
|
name: str | None = None,
|
|
):
|
|
self.cfg = self.Params.model_validate(params or {})
|
|
wren = _wren()
|
|
inputs = self._normalize_ports(requires)
|
|
outputs = self._normalize_ports(provides)
|
|
|
|
for spec in inputs:
|
|
allowed = INPUTS.get(spec.port)
|
|
if allowed is None:
|
|
raise ValueError(
|
|
f"'{spec.port}' is not an input of wren. It reads "
|
|
f"{', '.join(INPUTS)}."
|
|
)
|
|
if spec.dtype not in allowed:
|
|
raise ValueError(
|
|
f"'{spec.port}' is {spec.dtype.value}; wren reads it as "
|
|
f"{' or '.join(sorted(d.value for d in allowed))}."
|
|
)
|
|
if not any(spec.port == "prompt" for spec in inputs):
|
|
raise ValueError("wren needs a 'prompt' input: the command to act on.")
|
|
|
|
areas = {area.replace(" ", "_"): area for area in wren.AREAS}
|
|
#: (area, kind) to the ports standing for that device.
|
|
self.devices: dict[tuple[str, str], list[str]] = {}
|
|
for spec in outputs:
|
|
if spec.port == ACTIONS:
|
|
if spec.dtype is not DType.LIST or spec.item_dtype is not DType.RECORD:
|
|
raise ValueError(f"'{ACTIONS}' is a list of records.")
|
|
continue
|
|
device = _device(spec.port, areas)
|
|
if device is None:
|
|
raise ValueError(
|
|
f"'{spec.port}' names no device wren knows. A target is "
|
|
f"<area>_<kind>, the area one of {', '.join(areas)}, the kind "
|
|
f"one of {', '.join(KIND_WORDS)}."
|
|
)
|
|
if spec.dtype not in _TARGET_TYPES:
|
|
raise ValueError(
|
|
f"'{spec.port}' is {spec.dtype.value}; a target is bool, int, "
|
|
"float or str."
|
|
)
|
|
self.devices.setdefault(device, []).append(spec.port)
|
|
|
|
super().__init__(
|
|
self._act,
|
|
# The prompt is the event; the rest is read when it arrives.
|
|
requires=[
|
|
spec.model_copy(update={"trigger": spec.port == "prompt"})
|
|
for spec in inputs
|
|
],
|
|
provides=outputs,
|
|
params=params,
|
|
name=name or "wren",
|
|
)
|
|
|
|
def _house(
|
|
self,
|
|
wren: Any,
|
|
specs: dict[str, MessageSpec],
|
|
presence: str | None,
|
|
outdoor: float | None,
|
|
dark: bool | None,
|
|
) -> Any:
|
|
"""The house as the model reads it, each device in its current state."""
|
|
names = [specs[port].name for ports in self.devices.values() for port in ports]
|
|
current = (
|
|
self._pipeline.state.get_present(names)
|
|
if self._pipeline is not None
|
|
else {}
|
|
)
|
|
areas: dict[str, list[Any]] = {}
|
|
for area in wren.AREAS:
|
|
wired = {kind for (a, kind) in self.devices if a == area}
|
|
if not wired and area not in _TRAINED:
|
|
continue
|
|
kinds = wired | set(_TRAINED.get(area, ("light",)))
|
|
areas[area] = []
|
|
for kind in (k for k in wren.KINDS if k in kinds):
|
|
ports = self.devices.get((area, kind), [])
|
|
held = [
|
|
current[specs[p].name] for p in ports if specs[p].name in current
|
|
]
|
|
# The id is ours to choose: the model is shown aliases anyway.
|
|
areas[area].append(
|
|
wren.Device(kind, f"{area}:{kind}", _state(kind, held))
|
|
)
|
|
clock = time.localtime()
|
|
where = str(presence or "").strip().lower().replace("_", " ")
|
|
return wren.House(
|
|
time=time.strftime("%H:%M", clock),
|
|
# The range the model was trained on; past it a number is unfamiliar.
|
|
outdoor_c=min(max(float(20 if outdoor is None else outdoor), -5), 30),
|
|
dark=clock.tm_hour < 7 or clock.tm_hour >= 18 if dark is None else dark,
|
|
user_area=where if where in wren.AREAS else None,
|
|
areas=areas,
|
|
)
|
|
|
|
def _act(
|
|
self,
|
|
params: dict[str, Any],
|
|
prompt: str = "",
|
|
presence: str | None = None,
|
|
outdoor: float | None = None,
|
|
dark: bool | None = None,
|
|
) -> dict[str, Any] | None:
|
|
if not prompt.strip():
|
|
# An empty command is nothing to do, which is what a flow input
|
|
# starts as.
|
|
return None
|
|
wren = _wren()
|
|
specs = {spec.port: spec for spec in self.output_ports}
|
|
house = self._house(wren, specs, presence, outdoor, dark)
|
|
# One BLAS thread answers in about half the time for a model this
|
|
# small. Scoped to the call, so numpy elsewhere in the process keeps
|
|
# its threads.
|
|
with threadpool_limits(1):
|
|
calls, confidence = wren.act(prompt, house)
|
|
confidence = round(float(confidence), 3)
|
|
write = self.cfg.apply and confidence >= self.cfg.min_confidence
|
|
writes: dict[str, Any] = {}
|
|
actions: list[dict[str, Any]] = []
|
|
for call in calls:
|
|
area, _, kind = str(call.args.get("id", "")).partition(":")
|
|
ports = self.devices.get((area, kind), [])
|
|
if call.tool.split(".")[0] != _TOOLS.get(kind):
|
|
ports = []
|
|
hits = [
|
|
(specs[port], value)
|
|
for port in ports
|
|
if (value := _value(call, specs[port].dtype)) is not None
|
|
]
|
|
record = {
|
|
"message": "",
|
|
"value": None,
|
|
"applied": False,
|
|
"delay_s": call.delay_s,
|
|
"call": _text(call),
|
|
"confidence": confidence,
|
|
}
|
|
if not hits:
|
|
actions.append(record)
|
|
for spec, value in hits:
|
|
applied = write
|
|
if applied and call.delay_s:
|
|
applied = self._pipeline is not None and self._pipeline.defer(
|
|
self, self._to_messages({spec.port: value}) or {}, call.delay_s
|
|
)
|
|
elif applied:
|
|
writes[spec.port] = value
|
|
actions.append(
|
|
record | {"message": spec.name, "value": value, "applied": applied}
|
|
)
|
|
|
|
if any(spec.port == ACTIONS for spec in self.output_ports):
|
|
writes[ACTIONS] = actions
|
|
return writes or None
|