203 lines
7.7 KiB
Python
203 lines
7.7 KiB
Python
"""The AI Automation Suggester integration."""
|
|
import logging
|
|
import voluptuous as vol
|
|
|
|
from homeassistant.config_entries import ConfigEntry
|
|
from homeassistant.core import HomeAssistant, ServiceCall, callback
|
|
from homeassistant.exceptions import ConfigEntryNotReady, ServiceValidationError
|
|
from homeassistant.helpers.typing import ConfigType
|
|
|
|
from .api import async_register_http_views
|
|
from .const import (
|
|
DOMAIN,
|
|
PLATFORMS,
|
|
CONF_PROVIDER,
|
|
SERVICE_CLEAR_HISTORY,
|
|
SERVICE_GENERATE_SUGGESTIONS,
|
|
SERVICE_UPDATE_SUGGESTION,
|
|
ATTR_PROVIDER_CONFIG,
|
|
ATTR_CUSTOM_PROMPT,
|
|
CONFIG_VERSION
|
|
)
|
|
from .coordinator import AIAutomationCoordinator
|
|
from .store import async_get_suggestion_store
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
CONFIG_SCHEMA = vol.Schema({DOMAIN: vol.Schema({})}, extra=vol.ALLOW_EXTRA)
|
|
|
|
async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
|
|
"""Migrate old config entry if necessary."""
|
|
_LOGGER.debug(f"async_migrate_entry {config_entry.version}")
|
|
# Currently, no migration logic beyond ensuring version matches CONFIG_VERSION
|
|
if config_entry.version < CONFIG_VERSION:
|
|
_LOGGER.debug(f"Migrating config entry from version {config_entry.version} to {CONFIG_VERSION}")
|
|
new_data = {**config_entry.data}
|
|
new_data.pop('scan_frequency', None)
|
|
new_data.pop('initial_lag_time', None)
|
|
hass.config_entries.async_update_entry(config_entry, data=new_data, version=CONFIG_VERSION)
|
|
_LOGGER.debug("Migration successful")
|
|
return True
|
|
return True
|
|
|
|
|
|
def _listish(value):
|
|
"""Schema helper for service fields that can be a CSV string, list, or object."""
|
|
|
|
if value is None:
|
|
return []
|
|
if isinstance(value, (str, list, tuple, dict)):
|
|
return value
|
|
raise vol.Invalid("expected a list, comma-separated string, or object")
|
|
|
|
|
|
GENERATE_SUGGESTIONS_SCHEMA = vol.Schema(
|
|
{
|
|
vol.Optional(ATTR_PROVIDER_CONFIG): str,
|
|
vol.Optional(ATTR_CUSTOM_PROMPT): str,
|
|
vol.Optional("all_entities", default=False): bool,
|
|
vol.Optional("domains", default=[]): _listish,
|
|
vol.Optional("exclude_domains", default=[]): _listish,
|
|
vol.Optional("exclude_entities", default=[]): _listish,
|
|
vol.Optional("exclude_areas", default=[]): _listish,
|
|
vol.Optional("entity_limit", default=200): vol.All(vol.Coerce(int), vol.Range(min=1, max=2000)),
|
|
vol.Optional("automation_read_yaml", default=False): bool,
|
|
vol.Optional("automation_limit", default=100): vol.All(vol.Coerce(int), vol.Range(min=0, max=1000)),
|
|
}
|
|
)
|
|
|
|
|
|
UPDATE_SUGGESTION_SCHEMA = vol.Schema(
|
|
{
|
|
vol.Required("suggestion_id"): str,
|
|
vol.Required("status"): vol.In(["accepted", "declined", "dismissed", "new"]),
|
|
}
|
|
)
|
|
|
|
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
|
|
"""Set up the AI Automation Suggester component."""
|
|
hass.data.setdefault(DOMAIN, {})
|
|
async_register_http_views(hass)
|
|
|
|
async def handle_generate_suggestions(call: ServiceCall) -> None:
|
|
"""Handle the generate_suggestions service call."""
|
|
try:
|
|
coordinator = None
|
|
provider_config = call.data.get(ATTR_PROVIDER_CONFIG)
|
|
if provider_config:
|
|
coordinator = hass.data[DOMAIN].get(provider_config)
|
|
else:
|
|
# Find first available coordinator if none specified
|
|
for entry_id, coord in hass.data[DOMAIN].items():
|
|
if isinstance(coord, AIAutomationCoordinator):
|
|
coordinator = coord
|
|
break
|
|
|
|
if coordinator is None:
|
|
raise ServiceValidationError("No AI Automation Suggester provider configured")
|
|
|
|
await coordinator.async_generate_suggestions(
|
|
custom_prompt=call.data.get(ATTR_CUSTOM_PROMPT),
|
|
all_entities=call.data.get("all_entities", False),
|
|
domains=call.data.get("domains", []),
|
|
exclude_domains=call.data.get("exclude_domains", []),
|
|
exclude_entities=call.data.get("exclude_entities", []),
|
|
exclude_areas=call.data.get("exclude_areas", []),
|
|
entity_limit=call.data.get("entity_limit", 200),
|
|
automation_read_yaml=call.data.get("automation_read_yaml", False),
|
|
automation_limit=call.data.get("automation_limit", 100),
|
|
)
|
|
|
|
except KeyError:
|
|
raise ServiceValidationError("Provider configuration not found")
|
|
except Exception as err:
|
|
raise ServiceValidationError(f"Failed to generate suggestions: {err}")
|
|
|
|
async def handle_clear_history(call: ServiceCall) -> None:
|
|
"""Clear stored suggestion history."""
|
|
|
|
await async_get_suggestion_store(hass).async_clear()
|
|
|
|
async def handle_update_suggestion(call: ServiceCall) -> None:
|
|
"""Update a stored suggestion status."""
|
|
|
|
suggestion = await async_get_suggestion_store(hass).async_update_status(
|
|
call.data["suggestion_id"], call.data["status"]
|
|
)
|
|
if suggestion is None:
|
|
raise ServiceValidationError("Suggestion not found")
|
|
|
|
# Register the service
|
|
hass.services.async_register(
|
|
DOMAIN,
|
|
SERVICE_GENERATE_SUGGESTIONS,
|
|
handle_generate_suggestions,
|
|
schema=GENERATE_SUGGESTIONS_SCHEMA,
|
|
)
|
|
hass.services.async_register(
|
|
DOMAIN,
|
|
SERVICE_CLEAR_HISTORY,
|
|
handle_clear_history,
|
|
)
|
|
hass.services.async_register(
|
|
DOMAIN,
|
|
SERVICE_UPDATE_SUGGESTION,
|
|
handle_update_suggestion,
|
|
schema=UPDATE_SUGGESTION_SCHEMA,
|
|
)
|
|
|
|
return True
|
|
|
|
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
|
"""Set up AI Automation Suggester from a config entry."""
|
|
try:
|
|
if CONF_PROVIDER not in entry.data:
|
|
raise ConfigEntryNotReady("Provider not specified in config")
|
|
|
|
coordinator = AIAutomationCoordinator(hass, entry)
|
|
hass.data[DOMAIN][entry.entry_id] = coordinator
|
|
|
|
# Use the new async_forward_entry_setups method (plural) instead of the deprecated async_forward_entry_setup.
|
|
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
|
|
|
|
_LOGGER.debug(
|
|
"Setup complete for %s with provider %s",
|
|
entry.title,
|
|
entry.data.get(CONF_PROVIDER)
|
|
)
|
|
|
|
entry.async_on_unload(entry.add_update_listener(async_reload_entry))
|
|
|
|
@callback
|
|
def handle_custom_event(event):
|
|
_LOGGER.debug("Received custom event '%s', triggering suggestions with all_entities=True", event.event_type)
|
|
hass.async_create_task(coordinator_request_all_suggestions())
|
|
|
|
async def coordinator_request_all_suggestions():
|
|
await coordinator.async_generate_suggestions(all_entities=True)
|
|
|
|
entry.async_on_unload(hass.bus.async_listen("ai_automation_suggester_update", handle_custom_event))
|
|
|
|
return True
|
|
|
|
except Exception as err:
|
|
_LOGGER.error("Failed to setup integration: %s", err)
|
|
raise ConfigEntryNotReady from err
|
|
|
|
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
|
"""Unload a config entry."""
|
|
try:
|
|
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
|
|
if unload_ok:
|
|
coordinator = hass.data[DOMAIN].pop(entry.entry_id)
|
|
await coordinator.async_shutdown()
|
|
return unload_ok
|
|
except Exception as err:
|
|
_LOGGER.error("Error unloading entry: %s", err)
|
|
return False
|
|
|
|
async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
|
|
"""Reload config entry."""
|
|
await async_unload_entry(hass, entry)
|
|
await async_setup_entry(hass, entry)
|