Files
ha-text-ai/custom_components/ha_text_ai/config_flow.py
T
SMKRV 9f93f1ee18 Main changes:
Removed the validate_endpoint function
Optimized the validate_api_connection function
Simplified API connection verification
Preserved all error handling and retry logic
Improved exception handling
The integration should now correctly verify the
OpenAI API connection without false endpoint_not_available errors.
2024-11-19 17:33:27 +03:00

257 lines
8.3 KiB
Python

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