ha-text-ai/custom_components/ha_text_ai/config_flow.py

"""Config flow for HA text AI integration."""
from typing import Any, Dict, Optional, Tuple
import voluptuous as vol
import ssl
import certifi
import asyncio
from async_timeout import timeout
import aiohttp
from urllib.parse import urlparse

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 openai import AsyncOpenAI
from openai import OpenAIError, APIError, APIConnectionError, AuthenticationError, RateLimitError

from .const import (
    DOMAIN,
    CONF_MODEL,
    CONF_TEMPERATURE,
    CONF_MAX_TOKENS,
    CONF_API_ENDPOINT,
    CONF_REQUEST_INTERVAL,
    DEFAULT_MODEL,
    DEFAULT_TEMPERATURE,
    DEFAULT_MAX_TOKENS,
    DEFAULT_API_ENDPOINT,
    DEFAULT_REQUEST_INTERVAL,
)

import logging
_LOGGER = logging.getLogger(__name__)

# Create SSL context at module level
SSL_CONTEXT = ssl.create_default_context(cafile=certifi.where())

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=DEFAULT_TEMPERATURE
    ): vol.All(
        vol.Coerce(float),
        vol.Range(min=0, max=2)
    ),
    vol.Optional(
        CONF_MAX_TOKENS,
        default=DEFAULT_MAX_TOKENS
    ): vol.All(
        vol.Coerce(int),
        vol.Range(min=1, max=4096)
    ),
    vol.Optional(CONF_API_ENDPOINT, default=DEFAULT_API_ENDPOINT): str,
    vol.Optional(
        CONF_REQUEST_INTERVAL,
        default=DEFAULT_REQUEST_INTERVAL
    ): vol.All(
        vol.Coerce(float),
        vol.Range(min=0.1)
    ),
})

async def validate_api_connection(
    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."""
    for attempt in range(retry_count):
        try:
            async with timeout(10):
                client = AsyncOpenAI(
                    api_key=api_key,
                    base_url=endpoint,
                )

                models = await client.models.list()
                model_ids = [model.id for model in models.data]

                if model not in model_ids:
                    _LOGGER.warning(
                        "Model %s not found in available models: %s",
                        model,
                        ", ".join(model_ids)
                    )
                    return False, "invalid_model", model_ids
                return True, "", model_ids

        except asyncio.TimeoutError:
            _LOGGER.warning(
                "Timeout during API validation (attempt %d/%d)",
                attempt + 1,
                retry_count
            )
            if attempt == retry_count - 1:
                return False, "timeout", []
            await asyncio.sleep(retry_delay)

        except AuthenticationError as err:
            _LOGGER.error("Authentication error: %s", str(err))
            return False, "invalid_auth", []

        except RateLimitError as err:
            _LOGGER.error("Rate limit exceeded: %s", str(err))
            return False, "rate_limit", []

        except APIConnectionError as err:
            _LOGGER.error("API connection error: %s", str(err))
            return False, "cannot_connect", []

        except APIError as err:
            _LOGGER.error("API error: %s", str(err))
            return False, "api_error", []

        except Exception as err:
            _LOGGER.exception("Unexpected error during validation: %s", str(err))
            return False, "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
    ) -> Dict[str, Any]:
        """Handle the initial step."""
        errors: Dict[str, str] = {}

        if user_input is not None:
            try:
                # Validate URL format
                endpoint = user_input.get(CONF_API_ENDPOINT, DEFAULT_API_ENDPOINT)
                try:
                    result = urlparse(endpoint)
                    if not all([result.scheme, result.netloc]):
                        errors["base"] = "invalid_url_format"
                        return self.async_show_form(
                            step_id="user",
                            data_schema=STEP_USER_DATA_SCHEMA,
                            errors=errors
                        )
                except Exception as e:
                    _LOGGER.error("URL parsing error: %s", str(e))
                    errors["base"] = "invalid_url_format"
                    return self.async_show_form(
                        step_id="user",
                        data_schema=STEP_USER_DATA_SCHEMA,
                        errors=errors
                    )

                # Validate input data
                user_input = STEP_USER_DATA_SCHEMA(user_input)

                is_valid, error_code, available_models = await validate_api_connection(
                    user_input[CONF_API_KEY],
                    endpoint,
                    user_input[CONF_MODEL]
                )

                if is_valid:
                    await self.async_set_unique_id(user_input[CONF_API_KEY])
                    self._abort_if_unique_id_configured()

                    return self.async_create_entry(
                        title="HA text AI",
                        data=user_input
                    )

                errors["base"] = error_code
                if error_code == "invalid_model":
                    _LOGGER.warning(
                        "Selected model %s not found in available models: %s",
                        user_input[CONF_MODEL],
                        ", ".join(available_models)
                    )

            except vol.Invalid as err:
                _LOGGER.error("Validation error: %s", str(err))
                errors["base"] = "invalid_input"

        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,
            }
        )

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

class OptionsFlowHandler(config_entries.OptionsFlow):
    """Handle options flow for HA text AI."""

    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
    ) -> Dict[str, Any]:
        """Handle options flow."""
        if user_input is not None:
            return self.async_create_entry(title="", data=user_input)

        options_schema = vol.Schema({
            vol.Optional(
                CONF_TEMPERATURE,
                default=self.config_entry.options.get(
                    CONF_TEMPERATURE, DEFAULT_TEMPERATURE
                ),
                description={"suggested_value": DEFAULT_TEMPERATURE},
            ): vol.All(
                vol.Coerce(float),
                vol.Range(min=0, max=2)
            ),
            vol.Optional(
                CONF_MAX_TOKENS,
                default=self.config_entry.options.get(
                    CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS
                ),
                description={"suggested_value": DEFAULT_MAX_TOKENS},
            ): vol.All(
                vol.Coerce(int),
                vol.Range(min=1, max=4096)
            ),
            vol.Optional(
                CONF_REQUEST_INTERVAL,
                default=self.config_entry.options.get(
                    CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL
                ),
                description={"suggested_value": DEFAULT_REQUEST_INTERVAL},
            ): vol.All(
                vol.Coerce(float),
                vol.Range(min=0.1)
            ),
        })

        return self.async_show_form(
            step_id="init",
            data_schema=options_schema,
        )