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

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,
    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): cv.string,
        vol.Optional(CONF_MODEL, default=DEFAULT_MODEL): cv.string,
        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_API_ENDPOINT, default=DEFAULT_API_ENDPOINT): vol.All(
            cv.string,
            vol.Match(r'^https?://.+', msg="Must be a valid HTTP(S) URL")
        ),
        vol.Optional(CONF_REQUEST_INTERVAL, default=DEFAULT_REQUEST_INTERVAL): vol.All(
            vol.Coerce(float),
            vol.Range(min=MIN_REQUEST_INTERVAL)
        )
    }
)

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)

    # Determine API type and configure headers
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json"
    }

    # Configure endpoint and headers based on service type
    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"  # Using the correct v1/models endpoint
        _LOGGER.debug("Using VSE GPT endpoint: %s", models_url)

        # For VSE GPT, we might want to validate the model differently
        supported_models = [
            "anthropic/claude-3-5-haiku",
            "anthropic/claude-3.5-sonnet",
            # Add other supported VSE GPT models here
        ]

        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):
                # For VSE GPT, we'll skip the actual models endpoint check
                if "vsegpt" in endpoint.lower():
                    # Instead, let's verify the API key with a simple request
                    test_url = f"{base_url}/models"
                    async with session.get(test_url, headers=headers) as response:
                        if response.status == 404:
                            # This is actually expected for VSE GPT
                            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]

                # For other APIs, proceed with normal validation
                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:
        """Handle the initial step."""
        errors: Dict[str, str] = {}

        if user_input is not None:
            try:
                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
                    )

                validated_input = STEP_USER_DATA_SCHEMA(user_input)

                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 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
    ) -> FlowResult:
        """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
                ),
            ): 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)
            ),
        })

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