Files
ha-text-ai/custom_components/ha_text_ai/config_flow.py
T
SMKRV 58a7ae4229 Resoved blocking SSL verification issue in
coordinator.py

 API endpoint handling in config_flow.py
 changes

 Added support for the custom models in const.py

 Requirements in manifest.json updated
2024-11-20 12:07:30 +03:00

327 lines
11 KiB
Python

"""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
): cv.string,
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('/')
models_url = f"{base_url}/models"
_LOGGER.debug("Using VSE GPT endpoint: %s", models_url)
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):
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()
# Handle different API response formats
if "vsegpt" in endpoint.lower():
model_ids = [m["id"] for m in data.get("data", [])]
# Add VSE GPT specific model handling if needed
return True, "", model_ids
elif any(m in model.lower() for m in ["claude", "anthropic"]):
model_ids = [m["id"] for m in data.get("models", [])]
# Support for custom Anthropic model names
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,
)