diff --git a/custom_components/ha_text_ai/__init__.py b/custom_components/ha_text_ai/__init__.py index 8d94018..bf3ecee 100644 --- a/custom_components/ha_text_ai/__init__.py +++ b/custom_components/ha_text_ai/__init__.py @@ -4,17 +4,15 @@ from __future__ import annotations import logging import os import shutil -import asyncio -from datetime import datetime from datetime import datetime, timedelta -from typing import Any, Dict, Optional +from typing import Any, Dict import voluptuous as vol from async_timeout import timeout from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, CONF_NAME, Platform -from homeassistant.core import HomeAssistant, ServiceCall, callback +from homeassistant.core import HomeAssistant, ServiceCall from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError from homeassistant.helpers import config_validation as cv from homeassistant.helpers import aiohttp_client @@ -48,7 +46,6 @@ _LOGGER = logging.getLogger(__name__) CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) -# Обновленные схемы сервисов SERVICE_SCHEMA_ASK_QUESTION = vol.Schema({ vol.Required("instance"): cv.string, vol.Required("question"): cv.string, @@ -73,7 +70,6 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: """Set up the HA Text AI component.""" hass.data.setdefault(DOMAIN, {}) - # Copy custom icon try: source = os.path.join(os.path.dirname(__file__), 'icons', 'icon.svg') dest_dir = os.path.join(hass.config.path('www'), 'icons') @@ -87,10 +83,10 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: async def async_ask_question(call: ServiceCall) -> None: """Handle ask_question service.""" instance = call.data["instance"] - if instance not in hass.data[DOMAIN]: + coordinator = hass.data[DOMAIN].get(instance) + if not coordinator: raise HomeAssistantError(f"Instance {instance} not found") - coordinator = hass.data[DOMAIN][instance]["coordinator"] try: await coordinator.async_ask_question( question=call.data["question"], @@ -106,10 +102,10 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: async def async_clear_history(call: ServiceCall) -> None: """Handle clear_history service.""" instance = call.data["instance"] - if instance not in hass.data[DOMAIN]: + coordinator = hass.data[DOMAIN].get(instance) + if not coordinator: raise HomeAssistantError(f"Instance {instance} not found") - coordinator = hass.data[DOMAIN][instance]["coordinator"] try: await coordinator.async_clear_history() except Exception as err: @@ -119,10 +115,10 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: async def async_get_history(call: ServiceCall) -> None: """Handle get_history service.""" instance = call.data["instance"] - if instance not in hass.data[DOMAIN]: + coordinator = hass.data[DOMAIN].get(instance) + if not coordinator: raise HomeAssistantError(f"Instance {instance} not found") - coordinator = hass.data[DOMAIN][instance]["coordinator"] try: return await coordinator.async_get_history( limit=call.data.get("limit"), @@ -136,17 +132,16 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: async def async_set_system_prompt(call: ServiceCall) -> None: """Handle set_system_prompt service.""" instance = call.data["instance"] - if instance not in hass.data[DOMAIN]: + coordinator = hass.data[DOMAIN].get(instance) + if not coordinator: raise HomeAssistantError(f"Instance {instance} not found") - coordinator = hass.data[DOMAIN][instance]["coordinator"] try: await coordinator.async_set_system_prompt(call.data["prompt"]) except Exception as err: _LOGGER.error("Error setting system prompt: %s", str(err)) raise HomeAssistantError(f"Failed to set system prompt: {str(err)}") - # Register services hass.services.async_register( DOMAIN, SERVICE_ASK_QUESTION, @@ -208,19 +203,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: _LOGGER.error("API provider not specified") raise ConfigEntryNotReady("API provider is required") - # Убедимся, что домен инициализирован в hass.data - hass.data.setdefault(DOMAIN, {}) - session = aiohttp_client.async_get_clientsession(hass) api_provider = entry.data.get(CONF_API_PROVIDER) model = entry.data.get(CONF_MODEL, DEFAULT_MODEL) - endpoint = entry.data.get(CONF_API_ENDPOINT, + endpoint = entry.data.get( + CONF_API_ENDPOINT, DEFAULT_OPENAI_ENDPOINT if api_provider == API_PROVIDER_OPENAI - else DEFAULT_ANTHROPIC_ENDPOINT).rstrip('/') + else DEFAULT_ANTHROPIC_ENDPOINT + ).rstrip('/') api_key = entry.data[CONF_API_KEY] instance_name = entry.data.get(CONF_NAME, entry.entry_id) - is_anthropic = api_provider == API_PROVIDER_ANTHROPIC + headers = { "Content-Type": "application/json", "Accept": "application/json" @@ -235,23 +229,34 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if not await async_check_api(session, endpoint, headers, api_provider): raise ConfigEntryNotReady("API connection failed") + # Создаем API клиент + api_client = { + "session": session, + "endpoint": endpoint, + "headers": headers, + "api_provider": api_provider, + "model": model, + } + coordinator = HATextAICoordinator( hass=hass, - client=None, # Будет установлен позже + client=api_client, model=model, - update_interval=entry.data.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL), + update_interval=timedelta( + seconds=entry.data.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL) + ), instance_name=instance_name, max_tokens=entry.data.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS), temperature=entry.data.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE), is_anthropic=is_anthropic, ) - # Сохраняем данные интеграции используя entry.entry_id - hass.data[DOMAIN][entry.entry_id] = { - "coordinator": coordinator, - "config_entry": entry, - "instance_name": instance_name, # Сохраняем instance_name для доступа - } + # Инициализация координатора + await coordinator.async_config_entry_first_refresh() + + # Сохраняем координатор + hass.data.setdefault(DOMAIN, {}) + hass.data[DOMAIN][entry.entry_id] = coordinator # Загружаем платформы await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -265,7 +270,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" try: if entry.entry_id in hass.data[DOMAIN]: - coordinator = hass.data[DOMAIN][entry.entry_id]["coordinator"] + coordinator = hass.data[DOMAIN][entry.entry_id] await coordinator.async_shutdown() hass.data[DOMAIN].pop(entry.entry_id) diff --git a/custom_components/ha_text_ai/sensor.py b/custom_components/ha_text_ai/sensor.py index e602754..5b8599e 100644 --- a/custom_components/ha_text_ai/sensor.py +++ b/custom_components/ha_text_ai/sensor.py @@ -1,4 +1,4 @@ -"""Sensor platform for HA text AI.""" +"""Sensor platform for HA Text AI.""" import logging from typing import Any, Dict @@ -13,13 +13,9 @@ from homeassistant.util import dt as dt_util from homeassistant.util import slugify from .const import ( - # Основные константы домена DOMAIN, - CONF_INSTANCE, CONF_MODEL, CONF_API_PROVIDER, - - # Атрибуты сенсора ATTR_TOTAL_RESPONSES, ATTR_TOTAL_ERRORS, ATTR_AVG_RESPONSE_TIME, @@ -39,8 +35,6 @@ from .const import ( ATTR_API_STATUS, ATTR_RESPONSE, ATTR_QUESTION, - - # Метрики METRIC_TOTAL_TOKENS, METRIC_PROMPT_TOKENS, METRIC_COMPLETION_TOKENS, @@ -49,8 +43,6 @@ from .const import ( METRIC_AVERAGE_LATENCY, METRIC_MAX_LATENCY, METRIC_MIN_LATENCY, - - # Состояния и иконки STATE_READY, STATE_PROCESSING, STATE_ERROR, @@ -63,6 +55,8 @@ from .const import ( ENTITY_ICON_PROCESSING, ) +from .coordinator import HATextAICoordinator + _LOGGER = logging.getLogger(__name__) async def async_setup_entry( @@ -70,42 +64,55 @@ async def async_setup_entry( entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: - """Set up the HA text AI sensor.""" + """Set up the HA Text AI sensor.""" coordinator = hass.data[DOMAIN][entry.entry_id] async_add_entities([HATextAISensor(coordinator, entry)], True) class HATextAISensor(CoordinatorEntity, SensorEntity): - """HA text AI Sensor.""" + """HA Text AI Sensor.""" + + coordinator: HATextAICoordinator def __init__( self, - coordinator: dict, + coordinator: HATextAICoordinator, config_entry: ConfigEntry, ) -> None: """Initialize the sensor.""" super().__init__(coordinator) + self._config_entry = config_entry - self._name = config_entry.title - self._attr_unique_id = f"{config_entry.entry_id}_{slugify(self._name)}" - self.entity_id = f"sensor.ha_text_ai_{slugify(self._name)}" - self._attr_name = self._name + self._attr_name = config_entry.title + self._attr_unique_id = f"{config_entry.entry_id}_{slugify(self._attr_name)}" + self.entity_id = f"sensor.ha_text_ai_{slugify(self._attr_name)}" + self._current_state = STATE_INITIALIZING self._error_count = 0 self._last_error = None + self._last_update = None - # Получаем данные из конфигурации - model = self._config_entry.data.get(CONF_MODEL, "Unknown") - api_provider = self._config_entry.data.get(CONF_API_PROVIDER, "Unknown") + # Device info + model = config_entry.data.get(CONF_MODEL, "Unknown") + api_provider = config_entry.data.get(CONF_API_PROVIDER, "Unknown") - # Обновляем device_info с использованием DeviceInfo self._attr_device_info = DeviceInfo( identifiers={(DOMAIN, self._attr_unique_id)}, - name=self._name, + name=self._attr_name, manufacturer="Community", model=f"{model} ({api_provider} provider)", - sw_version="v1", + sw_version="1.0.0", ) + @property + def native_value(self) -> StateType: + """Return the native value of the sensor.""" + if self.coordinator.data is None: + return STATE_DISCONNECTED + + state = self.coordinator.data.get("state", STATE_READY) + self._current_state = state + return state + @property def icon(self) -> str: """Return the icon based on the current state.""" @@ -115,50 +122,21 @@ class HATextAISensor(CoordinatorEntity, SensorEntity): return ENTITY_ICON_PROCESSING return ENTITY_ICON - @property - def state(self) -> StateType: - """Return the state of the sensor.""" - try: - if isinstance(self.coordinator, dict): - last_update = self.coordinator.get("last_update") - if last_update: - if isinstance(last_update, str): - return dt_util.parse_datetime(last_update) - return last_update - except Exception as err: - _LOGGER.debug("Error parsing state: %s", err) - return self._current_state - @property def extra_state_attributes(self) -> Dict[str, Any]: """Return entity specific state attributes.""" attributes = { - ATTR_TOTAL_RESPONSES: 0, ATTR_MODEL: self._config_entry.data.get(CONF_MODEL, "Unknown"), + ATTR_API_PROVIDER: self._config_entry.data.get(CONF_API_PROVIDER, "Unknown"), ATTR_API_STATUS: self._current_state, - ATTR_ERROR_COUNT: self._error_count, + ATTR_TOTAL_ERRORS: self._error_count, ATTR_LAST_ERROR: self._last_error, } - if isinstance(self.coordinator, dict): - # Обновляем метрики если они есть - metrics = self.coordinator.get("metrics", {}) - if isinstance(metrics, dict): - attributes.update(metrics) + if self.coordinator.data: + data = self.coordinator.data - # Обновляем статус API - if "status" in self.coordinator: - attributes[ATTR_API_STATUS] = self.coordinator["status"] - - # Обновляем информацию о последнем ответе - last_response = self.coordinator.get("last_response", {}) - if isinstance(last_response, dict): - attributes.update({ - ATTR_RESPONSE: last_response.get("response", ""), - ATTR_QUESTION: last_response.get("question", ""), - }) - - # Обновляем дополнительные атрибуты + # Основные атрибуты for attr in [ ATTR_TOTAL_RESPONSES, ATTR_AVG_RESPONSE_TIME, @@ -171,50 +149,74 @@ class HATextAISensor(CoordinatorEntity, SensorEntity): ATTR_PERFORMANCE_METRICS, ATTR_HISTORY_SIZE, ATTR_UPTIME, + ATTR_SYSTEM_PROMPT, ]: - if attr in self.coordinator: - attributes[attr] = self.coordinator[attr] + if attr in data: + attributes[attr] = data[attr] + + # Метрики + if "metrics" in data: + metrics = data["metrics"] + for metric in [ + METRIC_TOTAL_TOKENS, + METRIC_PROMPT_TOKENS, + METRIC_COMPLETION_TOKENS, + METRIC_SUCCESSFUL_REQUESTS, + METRIC_FAILED_REQUESTS, + METRIC_AVERAGE_LATENCY, + METRIC_MAX_LATENCY, + METRIC_MIN_LATENCY, + ]: + if metric in metrics: + attributes[metric] = metrics[metric] + + # Последний ответ + if "last_response" in data: + last_response = data["last_response"] + attributes[ATTR_RESPONSE] = last_response.get("response", "") + attributes[ATTR_QUESTION] = last_response.get("question", "") return attributes - def _handle_coordinator_update(self) -> None: - """Handle updated data from the coordinator.""" - try: - if not isinstance(self.coordinator, dict): - self._current_state = STATE_DISCONNECTED - return - - # Определяем текущее состояние - if self.coordinator.get("is_processing"): - self._current_state = STATE_PROCESSING - elif self.coordinator.get("is_rate_limited"): - self._current_state = STATE_RATE_LIMITED - elif self.coordinator.get("is_maintenance"): - self._current_state = STATE_MAINTENANCE - else: - self._current_state = STATE_READY - - # Обновляем счетчик ошибок из метрик - metrics = self.coordinator.get("metrics", {}) - if isinstance(metrics, dict): - self._error_count = metrics.get("total_errors", self._error_count) - - except Exception as err: - _LOGGER.error("Error handling update: %s", err, exc_info=True) - self._error_count += 1 - self._last_error = str(err) - self._current_state = STATE_ERROR - - self.async_write_ha_state() - async def async_added_to_hass(self) -> None: """When entity is added to hass.""" await super().async_added_to_hass() self._handle_coordinator_update() - self._current_state = STATE_READY - async def async_reset_error_count(self) -> None: - """Reset the error counter.""" - self._error_count = 0 - self._last_error = None + def _handle_coordinator_update(self) -> None: + """Handle updated data from the coordinator.""" + if self.coordinator.data is None: + self._current_state = STATE_DISCONNECTED + self.async_write_ha_state() + return + + try: + data = self.coordinator.data + + # Обновляем состояние + if data.get("is_processing"): + self._current_state = STATE_PROCESSING + elif data.get("is_rate_limited"): + self._current_state = STATE_RATE_LIMITED + elif data.get("is_maintenance"): + self._current_state = STATE_MAINTENANCE + elif data.get("error"): + self._current_state = STATE_ERROR + self._last_error = data["error"] + self._error_count += 1 + else: + self._current_state = STATE_READY + + # Обновляем метрики + if "metrics" in data: + self._error_count = data["metrics"].get("total_errors", self._error_count) + + self._last_update = data.get("last_update") + + except Exception as err: + _LOGGER.error("Error handling update: %s", err, exc_info=True) + self._current_state = STATE_ERROR + self._last_error = str(err) + self._error_count += 1 + self.async_write_ha_state()