Release v2.0.0

custom_components/ha_text_ai/config_flow.py

@@ -1,14 +1,10 @@
 """Config flow for HA text AI integration."""
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional
 import voluptuous as vol
-import asyncio
 import aiohttp
-from async_timeout import timeout
-from urllib.parse import urlparse, urljoin
 
 from homeassistant import config_entries
 from homeassistant.const import CONF_API_KEY
-import homeassistant.helpers.config_validation as cv
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -20,403 +16,194 @@ from .const import (
     CONF_MAX_TOKENS,
     CONF_API_ENDPOINT,
     CONF_REQUEST_INTERVAL,
+    CONF_API_PROVIDER,
+    API_PROVIDER_OPENAI,
+    API_PROVIDER_ANTHROPIC,
+    API_PROVIDERS,
     DEFAULT_MODEL,
     DEFAULT_TEMPERATURE,
     DEFAULT_MAX_TOKENS,
-    DEFAULT_API_ENDPOINT,
     DEFAULT_REQUEST_INTERVAL,
+    DEFAULT_OPENAI_ENDPOINT,
+    DEFAULT_ANTHROPIC_ENDPOINT,
     MIN_TEMPERATURE,
     MAX_TEMPERATURE,
     MIN_MAX_TOKENS,
     MAX_MAX_TOKENS,
     MIN_REQUEST_INTERVAL,
-    API_VERSION,
-    API_MODELS_PATH,
-    ERROR_INVALID_API_KEY,
-    ERROR_CANNOT_CONNECT,
-    ERROR_UNKNOWN,
-    ERROR_INVALID_MODEL,
-    ERROR_RATE_LIMIT,
-    ERROR_API_ERROR,
-    ERROR_TIMEOUT,
 )
 
 import logging
 _LOGGER = logging.getLogger(__name__)
 
-STEP_USER_DATA_SCHEMA = vol.Schema({
-    vol.Required(CONF_API_KEY): str,
-    vol.Optional(CONF_MODEL, default=DEFAULT_MODEL): str,
-    vol.Optional(CONF_TEMPERATURE, default=str(DEFAULT_TEMPERATURE)): vol.All(
-        vol.Coerce(float),
-        vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
-    ),
-    vol.Optional(CONF_MAX_TOKENS, default=str(DEFAULT_MAX_TOKENS)): vol.All(
-        vol.Coerce(int),
-        vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
-    ),
-    vol.Optional(CONF_API_ENDPOINT, default=DEFAULT_API_ENDPOINT): str,
-    vol.Optional(CONF_REQUEST_INTERVAL, default=str(DEFAULT_REQUEST_INTERVAL)): vol.All(
-        vol.Coerce(float),
-        vol.Range(min=MIN_REQUEST_INTERVAL)
-    ),
-})
-
-async def async_step_user(
-    self,
-    user_input: Optional[Dict[str, Any]] = None
-) -> FlowResult:
-    """Handle the initial step."""
-    errors: Dict[str, str] = {}
-
-    if user_input is not None:
-        try:
-            _LOGGER.debug("Received user input: %s", user_input)
-
-            # Преобразование строковых значений в числовые
-            try:
-                temperature = float(user_input[CONF_TEMPERATURE])
-                max_tokens = int(user_input[CONF_MAX_TOKENS])
-                request_interval = float(user_input[CONF_REQUEST_INTERVAL])
-
-                # Обновляем значения в user_input
-                validated_input = {
-                    CONF_API_KEY: user_input[CONF_API_KEY],
-                    CONF_MODEL: user_input[CONF_MODEL],
-                    CONF_TEMPERATURE: temperature,
-                    CONF_MAX_TOKENS: max_tokens,
-                    CONF_API_ENDPOINT: user_input[CONF_API_ENDPOINT],
-                    CONF_REQUEST_INTERVAL: request_interval
-                }
-
-                _LOGGER.debug("Converted values: %s", validated_input)
-
-                # Проверка диапазонов
-                if not MIN_TEMPERATURE <= temperature <= MAX_TEMPERATURE:
-                    _LOGGER.error("Invalid temperature: %f", temperature)
-                    errors["base"] = "invalid_temperature"
-                elif not MIN_MAX_TOKENS <= max_tokens <= MAX_MAX_TOKENS:
-                    _LOGGER.error("Invalid max_tokens: %d", max_tokens)
-                    errors["base"] = "invalid_max_tokens"
-                elif request_interval < MIN_REQUEST_INTERVAL:
-                    _LOGGER.error("Invalid request_interval: %f", request_interval)
-                    errors["base"] = "invalid_request_interval"
-                else:
-                    # Проверка URL
-                    endpoint = validated_input[CONF_API_ENDPOINT]
-                    try:
-                        result = urlparse(endpoint)
-                        if not all([result.scheme, result.netloc]):
-                            errors["base"] = "invalid_url_format"
-                        else:
-                            is_valid, error_code, available_models = await validate_api_connection(
-                                self.hass,
-                                validated_input[CONF_API_KEY],
-                                endpoint,
-                                validated_input[CONF_MODEL]
-                            )
-
-                            if is_valid:
-                                await self.async_set_unique_id(validated_input[CONF_API_KEY])
-                                self._abort_if_unique_id_configured()
-                                return self.async_create_entry(
-                                    title="HA Text AI",
-                                    data=validated_input
-                                )
-                            errors["base"] = error_code
-                            if error_code == ERROR_INVALID_MODEL:
-                                _LOGGER.warning(
-                                    "Selected model %s not found in available models: %s",
-                                    validated_input[CONF_MODEL],
-                                    ", ".join(available_models)
-                                )
-                    except Exception as e:
-                        _LOGGER.error("URL parsing error: %s", str(e))
-                        errors["base"] = "invalid_url_format"
-
-            except ValueError as ve:
-                _LOGGER.error("Value conversion error: %s", str(ve))
-                errors["base"] = "invalid_input_format"
-
-        except Exception as err:
-            _LOGGER.error("Unexpected error: %s", str(err))
-            errors["base"] = "unknown"
-
-    return self.async_show_form(
-        step_id="user",
-        data_schema=STEP_USER_DATA_SCHEMA,
-        errors=errors,
-        description_placeholders={
-            "default_model": DEFAULT_MODEL,
-            "default_endpoint": DEFAULT_API_ENDPOINT,
-        }
-    )
-
-async def validate_api_connection(
-    hass,
-    api_key: str,
-    endpoint: str,
-    model: str,
-    retry_count: int = 3,
-    retry_delay: float = 1.0
-) -> Tuple[bool, str, list]:
-    """Validate API connection with retry logic."""
-    session = async_get_clientsession(hass)
-
-    headers = {
-        "Content-Type": "application/json",
-        "Accept": "application/json"
-    }
-
-    if "vsegpt" in endpoint.lower():
-        headers["Authorization"] = f"Bearer {api_key}"
-        base_url = endpoint.rstrip('/')
-        if not base_url.endswith("/v1"):
-            base_url = f"{base_url}/v1"
-        models_url = f"{base_url}/models"
-        _LOGGER.debug("Using VSE GPT endpoint: %s", models_url)
-
-        supported_models = [
-            "anthropic/claude-3-5-haiku",
-            "anthropic/claude-3.5-sonnet",
-        ]
-
-        if model in supported_models:
-            return True, "", supported_models
-
-    elif any(m in model.lower() for m in ["claude", "anthropic"]):
-        headers["x-api-key"] = api_key
-        headers["anthropic-version"] = "2023-06-01"
-        base_url = endpoint.rstrip('/')
-        if not base_url.endswith("/v1"):
-            base_url = f"{base_url}/v1"
-        models_url = f"{base_url}/models"
-        _LOGGER.debug("Using Anthropic endpoint: %s", models_url)
-    else:
-        headers["Authorization"] = f"Bearer {api_key}"
-        base_url = endpoint.rstrip('/')
-        if not base_url.endswith(f"/{API_VERSION}"):
-            base_url = f"{base_url}/{API_VERSION}"
-        models_url = f"{base_url}/{API_MODELS_PATH}"
-        _LOGGER.debug("Using OpenAI endpoint: %s", models_url)
-
-    for attempt in range(retry_count):
-        try:
-            async with timeout(10):
-                if "vsegpt" in endpoint.lower():
-                    test_url = f"{base_url}/models"
-                    async with session.get(test_url, headers=headers) as response:
-                        if response.status == 404:
-                            if model.startswith("anthropic/"):
-                                return True, "", [model]
-                        elif response.status == 401:
-                            return False, ERROR_INVALID_API_KEY, []
-                        elif response.status == 429:
-                            return False, ERROR_RATE_LIMIT, []
-                    return True, "", [model]
-
-                async with session.get(models_url, headers=headers) as response:
-                    _LOGGER.debug("API response status: %s", response.status)
-
-                    if response.status == 200:
-                        data = await response.json()
-
-                        if any(m in model.lower() for m in ["claude", "anthropic"]):
-                            model_ids = [m["id"] for m in data.get("models", [])]
-                            if model.startswith("anthropic/"):
-                                model = model.split("/")[1]
-                            if model in model_ids or any(m.endswith(model) for m in model_ids):
-                                return True, "", model_ids
-                        else:  # OpenAI format
-                            model_ids = [m["id"] for m in data.get("data", [])]
-                            if model in model_ids:
-                                return True, "", model_ids
-
-                        _LOGGER.warning(
-                            "Model %s not found in available models: %s",
-                            model,
-                            ", ".join(model_ids)
-                        )
-                        return False, ERROR_INVALID_MODEL, model_ids
-
-                    elif response.status == 401:
-                        _LOGGER.error("Authentication failed")
-                        return False, ERROR_INVALID_API_KEY, []
-
-                    elif response.status == 429:
-                        _LOGGER.error("Rate limit exceeded")
-                        if attempt < retry_count - 1:
-                            await asyncio.sleep(retry_delay * (2 ** attempt))
-                            continue
-                        return False, ERROR_RATE_LIMIT, []
-
-                    else:
-                        response_text = await response.text()
-                        _LOGGER.error(
-                            "API error: %s - %s",
-                            response.status,
-                            response_text
-                        )
-                        if attempt < retry_count - 1:
-                            await asyncio.sleep(retry_delay)
-                            continue
-                        return False, ERROR_API_ERROR, []
-
-        except asyncio.TimeoutError:
-            _LOGGER.warning(
-                "Timeout during API validation (attempt %d/%d)",
-                attempt + 1,
-                retry_count
-            )
-            if attempt < retry_count - 1:
-                await asyncio.sleep(retry_delay)
-                continue
-            return False, ERROR_TIMEOUT, []
-
-        except aiohttp.ClientError as err:
-            _LOGGER.error("Connection error: %s", str(err))
-            return False, ERROR_CANNOT_CONNECT, []
-
-        except Exception as err:
-            _LOGGER.exception("Unexpected error during validation: %s", str(err))
-            return False, ERROR_UNKNOWN, []
-
-    return False, ERROR_UNKNOWN, []
-
-
 class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for HA text AI."""
 
     VERSION = 1
 
-    async def async_step_user(
-        self,
-        user_input: Optional[Dict[str, Any]] = None
-    ) -> FlowResult:
+    async def async_step_user(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult:
         """Handle the initial step."""
-        errors: Dict[str, str] = {}
+        if user_input is None:
+            return self._show_provider_selection()
 
-        if user_input is not None:
-            try:
-                # Валидация значений через схему
-                validated_input = STEP_USER_DATA_SCHEMA(user_input)
+        if "provider" in user_input:
+            return self._show_provider_config(user_input["provider"])
 
-                endpoint = validated_input[CONF_API_ENDPOINT]
-                try:
-                    result = urlparse(endpoint)
-                    if not all([result.scheme, result.netloc]):
-                        errors["base"] = "invalid_url_format"
-                    else:
-                        is_valid, error_code, available_models = await validate_api_connection(
-                            self.hass,
-                            validated_input[CONF_API_KEY],
-                            endpoint,
-                            validated_input[CONF_MODEL]
-                        )
+        return await self._process_configuration(user_input)
 
-                        if is_valid:
-                            await self.async_set_unique_id(validated_input[CONF_API_KEY])
-                            self._abort_if_unique_id_configured()
-                            return self.async_create_entry(
-                                title="HA Text AI",
-                                data=validated_input
-                            )
-                        errors["base"] = error_code
-                        if error_code == ERROR_INVALID_MODEL:
-                            _LOGGER.warning(
-                                "Selected model %s not found in available models: %s",
-                                validated_input[CONF_MODEL],
-                                ", ".join(available_models)
-                            )
-                except Exception as e:
-                    _LOGGER.error("URL parsing error: %s", str(e))
-                    errors["base"] = "invalid_url_format"
+    def _show_provider_selection(self):
+        """Show provider selection screen."""
+        return self.async_show_form(
+            step_id="user",
+            data_schema=vol.Schema({
+                vol.Required("provider"): vol.In(API_PROVIDERS)
+            }),
+            description_placeholders={"providers": ", ".join(API_PROVIDERS)}
+        )
 
-            except vol.Invalid as err:
-                _LOGGER.error("Validation error: %s", str(err))
-                errors["base"] = "invalid_input"
-            except Exception as err:
-                _LOGGER.error("Unexpected error: %s", str(err))
-                errors["base"] = "unknown"
+    def _show_provider_config(self, provider):
+        """Show configuration screen for selected provider."""
+        default_endpoint = DEFAULT_OPENAI_ENDPOINT if provider == API_PROVIDER_OPENAI else DEFAULT_ANTHROPIC_ENDPOINT
 
         return self.async_show_form(
             step_id="user",
-            data_schema=STEP_USER_DATA_SCHEMA,
-            errors=errors,
-            description_placeholders={
-                "default_model": DEFAULT_MODEL,
-                "default_endpoint": DEFAULT_API_ENDPOINT,
-            }
+            data_schema=vol.Schema({
+                vol.Required(CONF_API_PROVIDER): vol.In([provider]),
+                vol.Required(CONF_API_KEY): str,
+                vol.Required(CONF_MODEL, default=DEFAULT_MODEL): str,
+                vol.Optional(CONF_API_ENDPOINT, default=default_endpoint): str,
+                vol.Optional(CONF_TEMPERATURE, default=DEFAULT_TEMPERATURE): vol.All(
+                    vol.Coerce(float),
+                    vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
+                ),
+                vol.Optional(CONF_MAX_TOKENS, default=DEFAULT_MAX_TOKENS): vol.All(
+                    vol.Coerce(int),
+                    vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
+                ),
+                vol.Optional(CONF_REQUEST_INTERVAL, default=DEFAULT_REQUEST_INTERVAL): vol.All(
+                    vol.Coerce(float),
+                    vol.Range(min=MIN_REQUEST_INTERVAL)
+                ),
+            }),
+            description_placeholders={"provider": provider}
         )
 
-    @staticmethod
-    @callback
-    def async_get_options_flow(
-        config_entry: config_entries.ConfigEntry,
-    ) -> config_entries.OptionsFlow:
-        """Get the options flow for this handler."""
-        return OptionsFlowHandler(config_entry)
+    async def _process_configuration(self, user_input):
+        """Validate and process user configuration."""
+        errors = {}
 
+        try:
+            session = async_get_clientsession(self.hass)
+
+            # Minimal API key validation
+            headers = {
+                "Authorization": f"Bearer {user_input[CONF_API_KEY]}",
+                "Content-Type": "application/json"
+            }
+
+            # Adjust headers and endpoint based on provider
+            if user_input[CONF_API_PROVIDER] == API_PROVIDER_ANTHROPIC:
+                headers = {
+                    "x-api-key": user_input[CONF_API_KEY],
+                    "anthropic-version": "2023-06-01"
+                }
+
+            # Basic connection test
+            try:
+                async with session.get(
+                    f"{user_input[CONF_API_ENDPOINT]}/models",
+                    headers=headers
+                ) as response:
+                    if response.status not in [200, 404]:
+                        errors["base"] = "cannot_connect"
+            except Exception as e:
+                _LOGGER.error(f"Connection test failed: {e}")
+                errors["base"] = "cannot_connect"
+
+            if not errors:
+                # Use a unique ID based on the API key
+                await self.async_set_unique_id(user_input[CONF_API_KEY])
+                self._abort_if_unique_id_configured()
+
+                return self.async_create_entry(
+                    title=f"HA Text AI ({user_input[CONF_API_PROVIDER]})",
+                    data=user_input
+                )
+
+        except Exception as e:
+            _LOGGER.error(f"Unexpected error: {e}")
+            errors["base"] = "unknown"
+
+        # Return to provider config, preserving previous input
+        return self.async_show_form(
+            step_id="user",
+            data_schema=vol.Schema({
+                vol.Required(CONF_API_PROVIDER): vol.In([user_input[CONF_API_PROVIDER]]),
+                vol.Required(CONF_API_KEY): str,
+                vol.Required(CONF_MODEL, default=user_input.get(CONF_MODEL, DEFAULT_MODEL)): str,
+                vol.Optional(CONF_API_ENDPOINT, default=user_input.get(CONF_API_ENDPOINT,
+                    DEFAULT_OPENAI_ENDPOINT if user_input[CONF_API_PROVIDER] == API_PROVIDER_OPENAI
+                    else DEFAULT_ANTHROPIC_ENDPOINT)): str,
+                vol.Optional(CONF_TEMPERATURE, default=user_input.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All(
+                    vol.Coerce(float),
+                    vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
+                ),
+                vol.Optional(CONF_MAX_TOKENS, default=user_input.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS)): vol.All(
+                    vol.Coerce(int),
+                    vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
+                ),
+                vol.Optional(CONF_REQUEST_INTERVAL, default=user_input.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All(
+                    vol.Coerce(float),
+                    vol.Range(min=MIN_REQUEST_INTERVAL)
+                ),
+            }),
+            description_placeholders={"provider": user_input[CONF_API_PROVIDER]},
+            errors=errors
+        )
+
+@staticmethod
+@callback
+def async_get_options_flow(config_entry: config_entries.ConfigEntry) -> config_entries.OptionsFlow:
+    """Create the options flow."""
+    return OptionsFlowHandler(config_entry)
 
 class OptionsFlowHandler(config_entries.OptionsFlow):
-    """Handle options flow for HA text AI."""
+    """Options flow handler."""
 
     def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
         """Initialize options flow."""
         self.config_entry = config_entry
 
-    async def async_step_init(
-        self,
-        user_input: Optional[Dict[str, Any]] = None
-    ) -> FlowResult:
-        """Handle options flow."""
+    async def async_step_init(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult:
+        """Manage the options."""
         if user_input is not None:
-            try:
-                # Validate using schema
-                options_schema = self._get_options_schema()
-                validated_input = options_schema(user_input)
-                return self.async_create_entry(title="", data=validated_input)
-            except vol.Invalid as err:
-                _LOGGER.error("Options validation error: %s", str(err))
-                return self.async_show_form(
-                    step_id="init",
-                    data_schema=self._get_options_schema(),
-                    errors={"base": "invalid_input"}
-                )
+            return self.async_create_entry(title="", data=user_input)
+
+        current_data = self.config_entry.data
 
         return self.async_show_form(
             step_id="init",
-            data_schema=self._get_options_schema(),
+            data_schema=vol.Schema({
+                vol.Optional(
+                    CONF_TEMPERATURE,
+                    default=current_data.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)
+                ): vol.All(
+                    vol.Coerce(float),
+                    vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
+                ),
+                vol.Optional(
+                    CONF_MAX_TOKENS,
+                    default=current_data.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS)
+                ): vol.All(
+                    vol.Coerce(int),
+                    vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
+                ),
+                vol.Optional(
+                    CONF_REQUEST_INTERVAL,
+                    default=current_data.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)
+                ): vol.All(
+                    vol.Coerce(float),
+                    vol.Range(min=MIN_REQUEST_INTERVAL)
+                ),
+            })
         )
-
-    def _get_options_schema(self) -> vol.Schema:
-        """Get options schema."""
-        return vol.Schema({
-            vol.Optional(
-                CONF_TEMPERATURE,
-                default=self.config_entry.options.get(
-                    CONF_TEMPERATURE, DEFAULT_TEMPERATURE
-                )
-            ): vol.All(
-                vol.Coerce(float),
-                vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
-            ),
-            vol.Optional(
-                CONF_MAX_TOKENS,
-                default=self.config_entry.options.get(
-                    CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS
-                )
-            ): vol.All(
-                vol.Coerce(int),
-                vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
-            ),
-            vol.Optional(
-                CONF_REQUEST_INTERVAL,
-                default=self.config_entry.options.get(
-                    CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL
-                )
-            ): vol.All(
-                vol.Coerce(float),
-                vol.Range(min=MIN_REQUEST_INTERVAL)
-            )
-        })