Release v2.0.0

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)
 

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()