From 7f463800546c187ceff8508de31e9beddf070a80 Mon Sep 17 00:00:00 2001 From: SMKRV Date: Mon, 25 Nov 2024 02:05:13 +0300 Subject: [PATCH] Release v2.0.0 --- custom_components/ha_text_ai/__init__.py | 57 +++++++++++++------- custom_components/ha_text_ai/api_client.py | 45 ++++------------ custom_components/ha_text_ai/config_flow.py | 59 +++++---------------- 3 files changed, 60 insertions(+), 101 deletions(-) diff --git a/custom_components/ha_text_ai/__init__.py b/custom_components/ha_text_ai/__init__.py index 1ef796c..e61708b 100644 --- a/custom_components/ha_text_ai/__init__.py +++ b/custom_components/ha_text_ai/__init__.py @@ -1,10 +1,4 @@ -# __init__.py -"""The HA Text AI integration. - -This module initializes the HA Text AI integration for Home Assistant, handling setup, configuration, -API connection verification, and managing service calls for interacting with AI models. -""" - +"""The HA Text AI integration.""" from __future__ import annotations import logging @@ -51,10 +45,8 @@ from .const import ( _LOGGER = logging.getLogger(__name__) -# Validate the configuration schema CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) -# Schema for the ask_question service SERVICE_SCHEMA_ASK_QUESTION = vol.Schema({ vol.Required("instance"): cv.string, vol.Required("question"): cv.string, @@ -64,13 +56,11 @@ SERVICE_SCHEMA_ASK_QUESTION = vol.Schema({ vol.Optional("max_tokens"): cv.positive_int, }) -# Schema for setting a system prompt SERVICE_SCHEMA_SET_SYSTEM_PROMPT = vol.Schema({ vol.Required("instance"): cv.string, vol.Required("prompt"): cv.string, }) -# Schema for getting history SERVICE_SCHEMA_GET_HISTORY = vol.Schema({ vol.Required("instance"): cv.string, vol.Optional("limit"): cv.positive_int, @@ -92,7 +82,6 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: """Set up the HA Text AI component.""" hass.data.setdefault(DOMAIN, {}) - # Try to copy custom icon to home assistant www directory try: source = os.path.join(os.path.dirname(__file__), 'icons', 'icon.svg') dest_dir = os.path.join(hass.config.path('www'), 'icons') @@ -148,7 +137,6 @@ async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: _LOGGER.error("Error setting system prompt: %s", str(err)) raise HomeAssistantError(f"Failed to set system prompt: {str(err)}") - # Registering services to handle different functionalities hass.services.async_register( DOMAIN, SERVICE_ASK_QUESTION, @@ -206,7 +194,6 @@ async def async_check_api(session, endpoint: str, headers: dict, provider: str) async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up HA Text AI from a config entry.""" try: - # Ensure the API provider is specified if CONF_API_PROVIDER not in entry.data: _LOGGER.error("API provider not specified") raise ConfigEntryNotReady("API provider is required") @@ -228,20 +215,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: "Accept": "application/json" } - # Set appropriate headers based on the provider if is_anthropic: headers["x-api-key"] = api_key headers["anthropic-version"] = "2023-06-01" else: headers["Authorization"] = f"Bearer {api_key}" - # Check the API connection if not await async_check_api(session, endpoint, headers, api_provider): raise ConfigEntryNotReady("API connection failed") _LOGGER.debug("Creating API client for %s with endpoint %s", api_provider, endpoint) - # Initialize the API client api_client = APIClient( session=session, endpoint=endpoint, @@ -250,7 +234,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: model=model, ) - # Initialize and set up the coordinator coordinator = HATextAICoordinator( hass=hass, client=api_client, @@ -265,4 +248,40 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: coordinator.data = coordinator._initial_state.copy() _LOGGER.debug(f"Initial state set for coordinator {instance_name}") - await coordinator.async_config_entry_first_refresh + 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) + + _LOGGER.info( + "Successfully set up %s instance '%s' with model %s", + api_provider, + instance_name, + model + ) + + return True + + except Exception as ex: + _LOGGER.exception("Setup error: %s", str(ex)) + raise ConfigEntryNotReady(f"Setup error: {str(ex)}") from ex + +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] + + if hasattr(coordinator.client, 'shutdown'): + await coordinator.client.shutdown() + + await coordinator.async_shutdown() + hass.data[DOMAIN].pop(entry.entry_id) + + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) + + except Exception as ex: + _LOGGER.exception("Error unloading entry: %s", str(ex)) + return False diff --git a/custom_components/ha_text_ai/api_client.py b/custom_components/ha_text_ai/api_client.py index bb73277..727af26 100644 --- a/custom_components/ha_text_ai/api_client.py +++ b/custom_components/ha_text_ai/api_client.py @@ -1,13 +1,7 @@ -# api_client.py -"""API Client for HA Text AI. - -This file contains the implementation of an API client for the HA Text AI integration compatible with both OpenAI and Anthropic service providers. -The client handles API requests, manages connection sessions, implements retry logic, and validates parameters for requests to ensure the correct functionality of the AI service integration within Home Assistant. -""" - +"""API Client for HA Text AI.""" import logging import asyncio -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from aiohttp import ClientSession, ClientTimeout from async_timeout import timeout @@ -26,10 +20,7 @@ from .const import ( _LOGGER = logging.getLogger(__name__) class APIClient: - """API Client for OpenAI and Anthropic. - - This client handles requests to OpenAI and Anthropic services to create text-based AI completions. - """ + """API Client for OpenAI and Anthropic.""" def __init__( self, @@ -39,10 +30,7 @@ class APIClient: api_provider: str, model: str, ) -> None: - """Initialize API client. - - Initialize with session, endpoint URL, request headers, type of API provider (OpenAI or Anthropic), and the AI model in use. - """ + """Initialize API client.""" self.session = session self.endpoint = endpoint self.headers = headers @@ -55,10 +43,7 @@ class APIClient: temperature: float, max_tokens: int, ) -> None: - """Validate API parameters. - - Check that the temperature and max_tokens are within their respective allowed ranges. - """ + """Validate API parameters.""" if not MIN_TEMPERATURE <= temperature <= MAX_TEMPERATURE: raise ValueError( f"Temperature must be between {MIN_TEMPERATURE} and {MAX_TEMPERATURE}" @@ -73,10 +58,7 @@ class APIClient: url: str, payload: Dict[str, Any], ) -> Dict[str, Any]: - """Make API request with retry logic. - - Attempt the API request a pre-defined number of times in case of failure, with an exponential backoff strategy for timeouts. - """ + """Make API request with retry logic.""" for attempt in range(API_RETRY_COUNT): try: async with timeout(API_TIMEOUT): @@ -107,10 +89,7 @@ class APIClient: temperature: float, max_tokens: int, ) -> Dict[str, Any]: - """Create completion using appropriate API. - - Make a completion request using either OpenAI or Anthropic API based on the provider specified during initialization. - """ + """Create completion using appropriate API.""" try: self._validate_parameters(temperature, max_tokens) @@ -133,10 +112,7 @@ class APIClient: temperature: float, max_tokens: int, ) -> Dict[str, Any]: - """Create completion using OpenAI API. - - Send details of the completion request to the OpenAI API and process the response. - """ + """Create completion using OpenAI API.""" url = f"{self.endpoint}/chat/completions" payload = { "model": model, @@ -168,10 +144,7 @@ class APIClient: temperature: float, max_tokens: int, ) -> Dict[str, Any]: - """Create completion using Anthropic API. - - Adapt messages format to Anthropic's requirement and make the API request to generate text completion. - """ + """Create completion using Anthropic API.""" url = f"{self.endpoint}/v1/messages" # Convert messages to Anthropic format diff --git a/custom_components/ha_text_ai/config_flow.py b/custom_components/ha_text_ai/config_flow.py index 61bfa18..f5a45c4 100644 --- a/custom_components/ha_text_ai/config_flow.py +++ b/custom_components/ha_text_ai/config_flow.py @@ -1,12 +1,4 @@ -# config_flow.py -"""Config flow for HA text AI integration. - -This module defines a configuration flow for the HA text AI integration in Home Assistant. -The flow guides the user through the setup process, including the selection of API provider, -configuration of API settings, and validation of the entered data. It ensures that the -configuration entries are unique and establish a valid connection to the AI API service. -""" - +"""Config flow for HA text AI integration.""" import logging from typing import Any, Dict, Optional @@ -56,10 +48,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._provider = None async def async_step_user(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult: - """Handle the initial step. - - This step asks the user to select an API provider. - """ + """Handle the initial step.""" if user_input is None: return self.async_show_form( step_id="user", @@ -77,11 +66,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return await self.async_step_provider() async def async_step_provider(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult: - """Handle provider configuration step. - - This step collects additional configuration details - based on the chosen provider like API key, model, etc. - """ + """Handle provider configuration step.""" if user_input is None: default_endpoint = ( DEFAULT_OPENAI_ENDPOINT if self._provider == API_PROVIDER_OPENAI @@ -113,21 +98,21 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): errors=self._errors ) + # Проверка уникальности имени instance_name = user_input[CONF_NAME] await self._async_validate_name(instance_name) if self._errors: return await self.async_step_provider() + # Проверка API подключения if not await self._async_validate_api(user_input): return await self.async_step_provider() + # Создание записи конфигурации return await self._create_entry(user_input) async def _async_validate_name(self, name: str) -> bool: - """Validate that the name is unique. - - Ensure no existing configuration entry has the same name. - """ + """Validate that the name is unique.""" for entry in self._async_current_entries(): if entry.data.get(CONF_NAME) == name: self._errors["name"] = "name_exists" @@ -135,10 +120,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return True async def _async_validate_api(self, user_input: Dict[str, Any]) -> bool: - """Validate API connection. - - Test the API connection to ensure the provided API details are correct. - """ + """Validate API connection.""" try: session = async_get_clientsession(self.hass) headers = self._get_api_headers(user_input) @@ -164,10 +146,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return False def _get_api_headers(self, user_input: Dict[str, Any]) -> Dict[str, str]: - """Get API headers based on provider. - - Returns appropriate headers for API requests depending on the provider. - """ + """Get API headers based on provider.""" api_key = user_input[CONF_API_KEY] if self._provider == API_PROVIDER_ANTHROPIC: @@ -182,10 +161,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): } async def _create_entry(self, user_input: Dict[str, Any]) -> FlowResult: - """Create the config entry. - - Constructs and adds the configuration entry to Home Assistant. - """ + """Create the config entry.""" instance_name = user_input[CONF_NAME] unique_id = f"{DOMAIN}_{instance_name}_{self._provider}".lower().replace(" ", "_") @@ -202,10 +178,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): @staticmethod @callback def async_get_options_flow(config_entry: config_entries.ConfigEntry) -> config_entries.OptionsFlow: - """Get the options flow for this handler. - - Provides access to configuration options after initial setup. - """ + """Get the options flow for this handler.""" return OptionsFlowHandler(config_entry) @@ -213,17 +186,11 @@ class OptionsFlowHandler(config_entries.OptionsFlow): """Handle options flow.""" def __init__(self, config_entry: config_entries.ConfigEntry) -> None: - """Initialize options flow. - - Prepare for managing configurable options of the integration. - """ + """Initialize options flow.""" self.config_entry = config_entry async def async_step_init(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult: - """Manage the options. - - Allows modification of configurable parameters post initial setup. - """ + """Manage the options.""" if user_input is not None: return self.async_create_entry(title="", data=user_input)