Files
ha-text-ai/custom_components/ha_text_ai/config_flow.py
T
2024-11-14 18:34:38 +03:00

124 lines
3.8 KiB
Python

"""Config flow for HA Text AI Integration."""
from __future__ import annotations
import logging
from typing import Any
import voluptuous as vol
import openai
from homeassistant import config_entries
from homeassistant.core import HomeAssistant, callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.exceptions import HomeAssistantError
from .const import (
DOMAIN,
CONF_API_KEY,
CONF_API_BASE,
CONF_REQUEST_INTERVAL,
DEFAULT_API_BASE,
DEFAULT_REQUEST_INTERVAL,
)
_LOGGER = logging.getLogger(__name__)
class TextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
"""Handle a config flow for HA Text AI Integration."""
VERSION = 1
async def async_step_user(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Handle the initial step."""
errors = {}
if user_input is not None:
try:
await self._test_api_key(user_input[CONF_API_KEY], user_input.get(CONF_API_BASE))
return self.async_create_entry(
title="HA Text AI",
data=user_input,
)
except ApiKeyError:
errors["base"] = "invalid_api_key"
except ApiConnectionError:
errors["base"] = "cannot_connect"
except Exception: # pylint: disable=broad-except
_LOGGER.exception("Unexpected exception")
errors["base"] = "unknown"
return self.async_show_form(
step_id="user",
data_schema=vol.Schema(
{
vol.Required(CONF_API_KEY): str,
vol.Optional(CONF_API_BASE, default=DEFAULT_API_BASE): str,
vol.Optional(CONF_REQUEST_INTERVAL, default=DEFAULT_REQUEST_INTERVAL): int,
}
),
errors=errors,
)
@staticmethod
async def _test_api_key(api_key: str, api_base: str | None) -> None:
"""Test if the API key is valid."""
try:
openai.api_key = api_key
if api_base:
openai.api_base = api_base
models = await openai.Model.alist()
if not models:
raise ApiKeyError
except openai.error.AuthenticationError as err:
raise ApiKeyError from err
except openai.error.APIConnectionError as err:
raise ApiConnectionError from err
@staticmethod
@callback
def async_get_options_flow(
config_entry: config_entries.ConfigEntry,
) -> TextAIOptionsFlow:
"""Get the options flow for this handler."""
return TextAIOptionsFlow(config_entry)
class TextAIOptionsFlow(config_entries.OptionsFlow):
"""Handle options flow for HA Text AI Integration."""
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: dict[str, Any] | None = None
) -> FlowResult:
"""Manage options."""
if user_input is not None:
return self.async_create_entry(title="", data=user_input)
return self.async_show_form(
step_id="init",
data_schema=vol.Schema(
{
vol.Optional(
CONF_REQUEST_INTERVAL,
default=self.config_entry.options.get(
CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL
),
): int,
}
),
)
class ApiKeyError(HomeAssistantError):
"""Error to indicate there is an invalid API key."""
class ApiConnectionError(HomeAssistantError):
"""Error to indicate there is a connection error."""