Add usage_prediction integration (#151206)

Co-authored-by: J. Nick Koston <[email protected]>
Co-authored-by: J. Nick Koston <[email protected]>

Author: balloob

CODEOWNERS

@@ -19,6 +19,7 @@
     "ssdp",
     "stream",
     "sun",
+    "usage_prediction",
     "usb",
     "webhook",
     "zeroconf"

homeassistant/components/default_config/manifest.json

@@ -0,0 +1,89 @@
+"""The usage prediction integration."""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from typing import Any
+
+from homeassistant.components import websocket_api
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.typing import ConfigType
+from homeassistant.util import dt as dt_util
+
+from . import common_control
+from .const import DATA_CACHE, DOMAIN
+from .models import EntityUsageDataCache, EntityUsagePredictions
+
+CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
+
+CACHE_DURATION = timedelta(hours=24)
+
+
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up the usage prediction integration."""
+    websocket_api.async_register_command(hass, ws_common_control)
+    hass.data[DATA_CACHE] = {}
+    return True
+
+
+@websocket_api.websocket_command({"type": f"{DOMAIN}/common_control"})
+@websocket_api.async_response
+async def ws_common_control(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
+    """Handle usage prediction common control WebSocket API."""
+    result = await get_cached_common_control(hass, connection.user.id)
+    time_category = common_control.time_category(dt_util.now().hour)
+    connection.send_result(
+        msg["id"],
+        {
+            "entities": getattr(result, time_category),
+        },
+    )
+
+
+async def get_cached_common_control(
+    hass: HomeAssistant, user_id: str
+) -> EntityUsagePredictions:
+    """Get cached common control predictions or fetch new ones.
+
+    Returns cached data if it's less than 24 hours old,
+    otherwise fetches new data and caches it.
+    """
+    # Create a unique storage key for this user
+    storage_key = user_id
+
+    cached_data = hass.data[DATA_CACHE].get(storage_key)
+
+    if isinstance(cached_data, asyncio.Task):
+        # If there's an ongoing task to fetch data, await its result
+        return await cached_data
+
+    # Check if cache is valid (less than 24 hours old)
+    if cached_data is not None:
+        if (dt_util.utcnow() - cached_data.timestamp) < CACHE_DURATION:
+            # Cache is still valid, return the cached predictions
+            return cached_data.predictions
+
+    # Create task fetching data
+    task = hass.async_create_task(
+        common_control.async_predict_common_control(hass, user_id)
+    )
+    hass.data[DATA_CACHE][storage_key] = task
+
+    try:
+        predictions = await task
+    except Exception:
+        # If the task fails, remove it from cache to allow retries
+        hass.data[DATA_CACHE].pop(storage_key)
+        raise
+
+    hass.data[DATA_CACHE][storage_key] = EntityUsageDataCache(
+        predictions=predictions,
+    )
+
+    return predictions

homeassistant/components/usage_prediction/__init__.py

@@ -0,0 +1,241 @@
+"""Code to generate common control usage patterns."""
+
+from __future__ import annotations
+
+from collections import Counter
+from collections.abc import Callable
+from datetime import datetime, timedelta
+from functools import cache
+import logging
+from typing import Any, Literal, cast
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from homeassistant.components.recorder import get_instance
+from homeassistant.components.recorder.db_schema import EventData, Events, EventTypes
+from homeassistant.components.recorder.models import uuid_hex_to_bytes_or_none
+from homeassistant.components.recorder.util import session_scope
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
+from homeassistant.util.json import json_loads_object
+
+from .models import EntityUsagePredictions
+
+_LOGGER = logging.getLogger(__name__)
+
+# Time categories for usage patterns
+TIME_CATEGORIES = ["morning", "afternoon", "evening", "night"]
+
+RESULTS_TO_INCLUDE = 8
+
+# List of domains for which we want to track usage
+ALLOWED_DOMAINS = {
+    # Entity platforms
+    Platform.AIR_QUALITY,
+    Platform.ALARM_CONTROL_PANEL,
+    Platform.BINARY_SENSOR,
+    Platform.BUTTON,
+    Platform.CALENDAR,
+    Platform.CAMERA,
+    Platform.CLIMATE,
+    Platform.COVER,
+    Platform.DATE,
+    Platform.DATETIME,
+    Platform.FAN,
+    Platform.HUMIDIFIER,
+    Platform.IMAGE,
+    Platform.LAWN_MOWER,
+    Platform.LIGHT,
+    Platform.LOCK,
+    Platform.MEDIA_PLAYER,
+    Platform.NUMBER,
+    Platform.SCENE,
+    Platform.SELECT,
+    Platform.SENSOR,
+    Platform.SIREN,
+    Platform.SWITCH,
+    Platform.TEXT,
+    Platform.TIME,
+    Platform.TODO,
+    Platform.UPDATE,
+    Platform.VACUUM,
+    Platform.VALVE,
+    Platform.WAKE_WORD,
+    Platform.WATER_HEATER,
+    Platform.WEATHER,
+    # Helpers with own domain
+    "counter",
+    "group",
+    "input_boolean",
+    "input_button",
+    "input_datetime",
+    "input_number",
+    "input_select",
+    "input_text",
+    "schedule",
+    "timer",
+}
+
+
+@cache
+def time_category(hour: int) -> Literal["morning", "afternoon", "evening", "night"]:
+    """Determine the time category for a given hour."""
+    if 6 <= hour < 12:
+        return "morning"
+    if 12 <= hour < 18:
+        return "afternoon"
+    if 18 <= hour < 22:
+        return "evening"
+    return "night"
+
+
+async def async_predict_common_control(
+    hass: HomeAssistant, user_id: str
+) -> EntityUsagePredictions:
+    """Generate a list of commonly used entities for a user.
+
+    Args:
+        hass: Home Assistant instance
+        user_id: User ID to filter events by.
+
+    Returns:
+        Dictionary with time categories as keys and lists of most common entity IDs as values
+    """
+    # Get the recorder instance to ensure it's ready
+    recorder = get_instance(hass)
+
+    # Execute the database operation in the recorder's executor
+    return await recorder.async_add_executor_job(
+        _fetch_with_session, hass, _fetch_and_process_data, user_id
+    )
+
+
+def _fetch_and_process_data(session: Session, user_id: str) -> EntityUsagePredictions:
+    """Fetch and process service call events from the database."""
+    # Prepare a dictionary to track results
+    results: dict[str, Counter[str]] = {
+        time_cat: Counter() for time_cat in TIME_CATEGORIES
+    }
+
+    # Keep track of contexts that we processed so that we will only process
+    # the first service call in a context, and not subsequent calls.
+    context_processed: set[bytes] = set()
+    thirty_days_ago_ts = (dt_util.utcnow() - timedelta(days=30)).timestamp()
+    user_id_bytes = uuid_hex_to_bytes_or_none(user_id)
+    if not user_id_bytes:
+        raise ValueError("Invalid user_id format")
+
+    # Build the main query for events with their data
+    query = (
+        select(
+            Events.context_id_bin,
+            Events.time_fired_ts,
+            EventData.shared_data,
+        )
+        .select_from(Events)
+        .outerjoin(EventData, Events.data_id == EventData.data_id)
+        .outerjoin(EventTypes, Events.event_type_id == EventTypes.event_type_id)
+        .where(Events.time_fired_ts >= thirty_days_ago_ts)
+        .where(Events.context_user_id_bin == user_id_bytes)
+        .where(EventTypes.event_type == "call_service")
+        .order_by(Events.time_fired_ts)
+    )
+
+    # Execute the query
+    context_id: bytes
+    time_fired_ts: float
+    shared_data: str | None
+    local_time_zone = dt_util.get_default_time_zone()
+    for context_id, time_fired_ts, shared_data in (
+        session.connection().execute(query).all()
+    ):
+        # Skip if we have already processed an event that was part of this context
+        if context_id in context_processed:
+            continue
+
+        # Mark this context as processed
+        context_processed.add(context_id)
+
+        # Parse the event data
+        if not shared_data:
+            continue
+
+        try:
+            event_data = json_loads_object(shared_data)
+        except (ValueError, TypeError) as err:
+            _LOGGER.debug("Failed to parse event data: %s", err)
+            continue
+
+        # Empty event data, skipping
+        if not event_data:
+            continue
+
+        service_data = cast(dict[str, Any] | None, event_data.get("service_data"))
+
+        # No service data found, skipping
+        if not service_data:
+            continue
+
+        entity_ids: str | list[str] | None
+        if (target := service_data.get("target")) and (
+            target_entity_ids := target.get("entity_id")
+        ):
+            entity_ids = target_entity_ids
+        else:
+            entity_ids = service_data.get("entity_id")
+
+        # No entity IDs found, skip this event
+        if entity_ids is None:
+            continue
+
+        if not isinstance(entity_ids, list):
+            entity_ids = [entity_ids]
+
+        # Filter out entity IDs that are not in allowed domains
+        entity_ids = [
+            entity_id
+            for entity_id in entity_ids
+            if entity_id.split(".")[0] in ALLOWED_DOMAINS
+        ]
+
+        if not entity_ids:
+            continue
+
+        # Convert timestamp to datetime and determine time category
+        if time_fired_ts:
+            # Convert to local time for time category determination
+            period = time_category(
+                datetime.fromtimestamp(time_fired_ts, local_time_zone).hour
+            )
+
+            # Count entity usage
+            for entity_id in entity_ids:
+                results[period][entity_id] += 1
+
+    return EntityUsagePredictions(
+        morning=[
+            ent_id for (ent_id, _) in results["morning"].most_common(RESULTS_TO_INCLUDE)
+        ],
+        afternoon=[
+            ent_id
+            for (ent_id, _) in results["afternoon"].most_common(RESULTS_TO_INCLUDE)
+        ],
+        evening=[
+            ent_id for (ent_id, _) in results["evening"].most_common(RESULTS_TO_INCLUDE)
+        ],
+        night=[
+            ent_id for (ent_id, _) in results["night"].most_common(RESULTS_TO_INCLUDE)
+        ],
+    )
+
+
+def _fetch_with_session(
+    hass: HomeAssistant,
+    fetch_func: Callable[[Session], EntityUsagePredictions],
+    *args: object,
+) -> EntityUsagePredictions:
+    """Execute a fetch function with a database session."""
+    with session_scope(hass=hass, read_only=True) as session:
+        return fetch_func(session, *args)

homeassistant/components/usage_prediction/common_control.py

@@ -0,0 +1,13 @@
+"""Constants for the usage prediction integration."""
+
+import asyncio
+
+from homeassistant.util.hass_dict import HassKey
+
+from .models import EntityUsageDataCache, EntityUsagePredictions
+
+DOMAIN = "usage_prediction"
+
+DATA_CACHE: HassKey[
+    dict[str, asyncio.Task[EntityUsagePredictions] | EntityUsageDataCache]
+] = HassKey("usage_prediction")

homeassistant/components/usage_prediction/const.py

@@ -0,0 +1,10 @@
+{
+  "domain": "usage_prediction",
+  "name": "Usage Prediction",
+  "codeowners": ["@home-assistant/core"],
+  "dependencies": ["http", "recorder"],
+  "documentation": "https://www.home-assistant.io/integrations/usage_prediction",
+  "integration_type": "system",
+  "iot_class": "calculated",
+  "quality_scale": "internal"
+}

homeassistant/components/usage_prediction/manifest.json

@@ -0,0 +1,24 @@
+"""Models for the usage prediction integration."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+
+from homeassistant.util import dt as dt_util
+
+
+@dataclass
+class EntityUsagePredictions:
+    """Prediction which entities are likely to be used in each time category."""
+
+    morning: list[str] = field(default_factory=list)
+    afternoon: list[str] = field(default_factory=list)
+    evening: list[str] = field(default_factory=list)
+    night: list[str] = field(default_factory=list)
+
+
+@dataclass
+class EntityUsageDataCache:
+    """Data model for entity usage prediction."""
+
+    predictions: EntityUsagePredictions
+    timestamp: datetime = field(default_factory=dt_util.utcnow)