Initial Home Assistant commit
This commit is contained in:
@@ -0,0 +1,202 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user