diff --git a/custom_components/ha_text_ai/__init__.py b/custom_components/ha_text_ai/__init__.py index 2e792bd..1a0a56e 100644 --- a/custom_components/ha_text_ai/__init__.py +++ b/custom_components/ha_text_ai/__init__.py @@ -9,7 +9,7 @@ from datetime import datetime from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, ServiceCall from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import config_validation as cv from homeassistant.helpers import aiohttp_client @@ -34,6 +34,11 @@ from .const import ( API_TIMEOUT, API_RETRY_COUNT, API_BACKOFF_FACTOR, + SERVICE_ASK_QUESTION, + SERVICE_CLEAR_HISTORY, + SERVICE_SET_SYSTEM_PROMPT, + SERVICE_SCHEMA_ASK_QUESTION, + SERVICE_SCHEMA_SET_SYSTEM_PROMPT, ) _LOGGER = logging.getLogger(__name__) @@ -133,46 +138,70 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # Initialize the coordinator await coordinator.async_initialize() - - # Perform first refresh await coordinator.async_config_entry_first_refresh() - # Check coordinator status - if coordinator.endpoint_status == "auth_error": - raise ConfigEntryNotReady("Authentication failed") - elif coordinator.endpoint_status == "rate_limited": - _LOGGER.warning("API rate limited during setup") - elif coordinator.endpoint_status == "maintenance": - raise ConfigEntryNotReady("API is in maintenance mode") - elif coordinator.endpoint_status == "error": - raise ConfigEntryNotReady("API error during setup") - elif not coordinator.last_update_success: - raise ConfigEntryNotReady("Failed to initialize coordinator") - hass.data[DOMAIN][entry.entry_id] = coordinator + # Register services + async def async_ask_question(call: ServiceCall) -> None: + """Handle ask_question service.""" + question = call.data.get("question", "") + if not question: + _LOGGER.error("No question provided") + return + + params = { + "system_prompt": call.data.get("system_prompt"), + "model": call.data.get("model"), + "temperature": call.data.get("temperature"), + "max_tokens": call.data.get("max_tokens"), + "priority": call.data.get("priority", False) + } + + try: + await coordinator.async_ask_question(question, **params) + except Exception as err: + _LOGGER.error("Error asking question: %s", str(err)) + + async def async_clear_history(call: ServiceCall) -> None: + """Handle clear_history service.""" + try: + await coordinator.clear_history() + except Exception as err: + _LOGGER.error("Error clearing history: %s", str(err)) + + async def async_set_system_prompt(call: ServiceCall) -> None: + """Handle set_system_prompt service.""" + prompt = call.data.get("prompt", "") + if prompt: + coordinator.update_system_prompt(prompt) + else: + _LOGGER.error("No prompt provided") + + # Register all services + hass.services.async_register( + DOMAIN, + SERVICE_ASK_QUESTION, + async_ask_question, + schema=SERVICE_SCHEMA_ASK_QUESTION + ) + + hass.services.async_register( + DOMAIN, + SERVICE_CLEAR_HISTORY, + async_clear_history + ) + + hass.services.async_register( + DOMAIN, + SERVICE_SET_SYSTEM_PROMPT, + async_set_system_prompt, + schema=SERVICE_SCHEMA_SET_SYSTEM_PROMPT + ) + # Set up platforms await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True - except Exception as ex: - _LOGGER.exception("Setup error: %s", str(ex)) - raise ConfigEntryNotReady from ex - -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: - """Unload a config entry.""" - try: - coordinator = hass.data[DOMAIN].get(entry.entry_id) - if coordinator: - await coordinator.async_shutdown() - - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok - - except Exception as ex: - _LOGGER.exception("Error unloading entry: %s", str(ex)) - return False + except Exception