diff --git a/custom_components/ha_text_ai/api_client.py b/custom_components/ha_text_ai/api_client.py index a0931c7..733ff78 100644 --- a/custom_components/ha_text_ai/api_client.py +++ b/custom_components/ha_text_ai/api_client.py @@ -429,7 +429,6 @@ class APIClient: def generate_content(): # For single message without history, use generate_content if len(contents) <= 1: - # If we have no content yet, create a simple prompt if not contents: prompt = "I need your assistance." else: @@ -441,16 +440,30 @@ class APIClient: config=config ) else: - # For multi-turn conversations, use chat - chat = client.chats.create(model=model, config=config) + # For multi-turn conversations, pass history to chat + # and only send the last user message + last_user_msg = None + history = [] - # Send all messages in sequence - for content in contents: - if content["role"] == "user": - response = chat.send_message(content["parts"][0]["text"]) - # We don't send assistant messages as they're already part of the history + # Find the last user message — that's the new query + for i in range(len(contents) - 1, -1, -1): + if contents[i]["role"] == "user": + last_user_msg = contents[i]["parts"][0]["text"] + history = contents[:i] + break - return response + if last_user_msg is None: + # No user messages at all — shouldn't happen, but handle gracefully + return client.models.generate_content( + model=model, + contents="I need your assistance.", + config=config + ) + + chat = client.chats.create( + model=model, config=config, history=history + ) + return chat.send_message(last_user_msg) response = await asyncio.to_thread(generate_content) diff --git a/custom_components/ha_text_ai/config_flow.py b/custom_components/ha_text_ai/config_flow.py index 1ff6916..4cbad01 100644 --- a/custom_components/ha_text_ai/config_flow.py +++ b/custom_components/ha_text_ai/config_flow.py @@ -8,7 +8,6 @@ Config flow for HA text AI integration. """ import logging from typing import Any, Dict, Optional -from datetime import datetime, timedelta import voluptuous as vol from homeassistant import config_entries @@ -49,8 +48,10 @@ from .const import ( DEFAULT_MAX_HISTORY, CONF_MAX_HISTORY_SIZE, ) -from .utils import normalize_name, safe_log_data, validate_endpoint # noqa: F401 — re-exported for backward compat -from .providers import get_default_endpoint, get_default_model +from homeassistant.util import dt as dt_util + +from .utils import normalize_name, safe_log_data, validate_endpoint +from .providers import get_default_endpoint, get_default_model, build_auth_headers _LOGGER = logging.getLogger(__name__) @@ -84,52 +85,56 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._provider = user_input[CONF_API_PROVIDER] return await self.async_step_provider() + def _build_provider_schema( + self, data: Optional[Dict[str, Any]] = None + ) -> vol.Schema: + """Build provider configuration schema with optional defaults from data.""" + defaults = data or {} + return vol.Schema({ + vol.Required(CONF_NAME, default=defaults.get(CONF_NAME, "my_assistant")): str, + vol.Required(CONF_API_KEY): str, + vol.Required(CONF_MODEL, default=defaults.get(CONF_MODEL, get_default_model(self._provider))): str, + vol.Required(CONF_API_ENDPOINT, default=defaults.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str, + vol.Optional(CONF_TEMPERATURE, default=defaults.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All( + vol.Coerce(float), + vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE) + ), + vol.Optional(CONF_MAX_TOKENS, default=defaults.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=defaults.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All( + vol.Coerce(float), + vol.Range(min=MIN_REQUEST_INTERVAL) + ), + vol.Optional(CONF_API_TIMEOUT, default=defaults.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All( + vol.Coerce(int), + vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT) + ), + vol.Optional( + CONF_CONTEXT_MESSAGES, + default=defaults.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES) + ): vol.All( + vol.Coerce(int), + vol.Range(min=1, max=20) + ), + vol.Optional( + CONF_MAX_HISTORY_SIZE, + default=defaults.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY) + ): vol.All( + vol.Coerce(int), + vol.Range(min=1, max=100) + ), + }) + async def async_step_provider(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult: """Handle provider configuration step.""" self._errors = {} if user_input is None: - default_endpoint = get_default_endpoint(self._provider) - default_model = get_default_model(self._provider) - return self.async_show_form( step_id="provider", - data_schema=vol.Schema({ - vol.Required(CONF_NAME, default="my_assistant"): str, - vol.Required(CONF_API_KEY): str, - vol.Required(CONF_MODEL, default=default_model): str, - vol.Required(CONF_API_ENDPOINT, default=default_endpoint): str, - 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_REQUEST_INTERVAL, default=DEFAULT_REQUEST_INTERVAL): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_REQUEST_INTERVAL) - ), - vol.Optional(CONF_API_TIMEOUT, default=DEFAULT_API_TIMEOUT): vol.All( - vol.Coerce(int), - vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT) - ), - vol.Optional( - CONF_CONTEXT_MESSAGES, - default=DEFAULT_CONTEXT_MESSAGES - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=20) - ), - vol.Optional( - CONF_MAX_HISTORY_SIZE, - default=DEFAULT_MAX_HISTORY - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=100) - ), - }) + data_schema=self._build_provider_schema(), ) _LOGGER.debug("Provider step input data: %s", safe_log_data(user_input)) @@ -139,7 +144,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): # Check if CONF_NAME exists in input_copy and ensure it's not empty if CONF_NAME not in input_copy or not input_copy[CONF_NAME]: _LOGGER.warning("Missing name in configuration input: %s", safe_log_data(input_copy)) - input_copy[CONF_NAME] = f"assistant_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + input_copy[CONF_NAME] = f"assistant_{dt_util.utcnow().strftime('%Y%m%d_%H%M%S')}" _LOGGER.info("Auto-generated name: %s", input_copy[CONF_NAME]) # Ensure API key is present @@ -148,168 +153,35 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): _LOGGER.error("API validation error: 'api_key'") return self.async_show_form( step_id="provider", - data_schema=vol.Schema({ - vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str, - vol.Required(CONF_API_KEY): str, - vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str, - vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str, - vol.Optional(CONF_TEMPERATURE, default=input_copy.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE) - ), - vol.Optional(CONF_MAX_TOKENS, default=input_copy.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=input_copy.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_REQUEST_INTERVAL) - ), - vol.Optional(CONF_API_TIMEOUT, default=input_copy.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All( - vol.Coerce(int), - vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT) - ), - vol.Optional( - CONF_CONTEXT_MESSAGES, - default=input_copy.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES) - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=20) - ), - vol.Optional( - CONF_MAX_HISTORY_SIZE, - default=input_copy.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY) - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=100) - ), - }), + data_schema=self._build_provider_schema(input_copy), errors=self._errors ) try: - # Validate and normalize the name normalized_name = self._validate_and_normalize_name(input_copy[CONF_NAME]) input_copy[CONF_NAME] = normalized_name except ValueError as e: return self.async_show_form( step_id="provider", - data_schema=vol.Schema({ - vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str, - vol.Required(CONF_API_KEY): str, - vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str, - vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str, - vol.Optional(CONF_TEMPERATURE, default=input_copy.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE) - ), - vol.Optional(CONF_MAX_TOKENS, default=input_copy.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=input_copy.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_REQUEST_INTERVAL) - ), - vol.Optional(CONF_API_TIMEOUT, default=input_copy.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All( - vol.Coerce(int), - vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT) - ), - vol.Optional( - CONF_CONTEXT_MESSAGES, - default=input_copy.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES) - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=20) - ), - vol.Optional( - CONF_MAX_HISTORY_SIZE, - default=input_copy.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY) - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=100) - ), - }), + data_schema=self._build_provider_schema(input_copy), errors={"name": str(e)} ) try: - # Special handling for Gemini API validation - if self._provider == API_PROVIDER_GEMINI: - # For Gemini, we just check if API key is present as there's no simple endpoint to validate - if not input_copy.get(CONF_API_KEY): - self._errors["base"] = "invalid_auth" - _LOGGER.error("API validation error: 'api_key'") - return self.async_show_form( - step_id="provider", - data_schema=vol.Schema({ - vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str, - vol.Required(CONF_API_KEY): str, - vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str, - vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str, - # Other fields remain the same - }), - errors=self._errors - ) - else: - # For other providers, validate API connection - if not await self._async_validate_api(input_copy): - return self.async_show_form( - step_id="provider", - data_schema=vol.Schema({ - vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str, - vol.Required(CONF_API_KEY): str, - vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str, - vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str, - vol.Optional(CONF_TEMPERATURE, default=input_copy.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE) - ), - vol.Optional(CONF_MAX_TOKENS, default=input_copy.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=input_copy.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All( - vol.Coerce(float), - vol.Range(min=MIN_REQUEST_INTERVAL) - ), - vol.Optional(CONF_API_TIMEOUT, default=input_copy.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All( - vol.Coerce(int), - vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT) - ), - vol.Optional( - CONF_CONTEXT_MESSAGES, - default=input_copy.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES) - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=20) - ), - vol.Optional( - CONF_MAX_HISTORY_SIZE, - default=input_copy.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY) - ): vol.All( - vol.Coerce(int), - vol.Range(min=1, max=100) - ), - }), - errors=self._errors - ) - except Exception as e: - # Handle any unexpected exceptions during validation + if not await self._async_validate_api(input_copy): + return self.async_show_form( + step_id="provider", + data_schema=self._build_provider_schema(input_copy), + errors=self._errors + ) + except Exception: _LOGGER.exception("Unexpected error during API validation") return self.async_show_form( step_id="provider", - data_schema=vol.Schema({ - vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str, - vol.Required(CONF_API_KEY): str, - vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str, - vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str, - # Other fields remain the same - }), + data_schema=self._build_provider_schema(input_copy), errors={"base": "unknown"} ) - # All validation passed, create the entry return await self._create_entry(input_copy) def _validate_and_normalize_name(self, name: str) -> str: @@ -345,14 +217,13 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return normalized async def _async_validate_api(self, user_input: Dict[str, Any]) -> bool: - """Validate API connection.""" + """Validate API connection using provider registry.""" try: if CONF_API_KEY not in user_input: _LOGGER.error("API validation error: 'api_key'") self._errors["base"] = "invalid_auth" return False - # Validate endpoint URL (HTTPS-only, SSRF protection) try: endpoint = validate_endpoint(user_input[CONF_API_ENDPOINT]) except ValueError as err: @@ -360,57 +231,33 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._errors["base"] = "cannot_connect" return False - session = async_get_clientsession(self.hass) - headers = self._get_api_headers(user_input) - if self._provider == API_PROVIDER_GEMINI: if not user_input[CONF_API_KEY]: self._errors["base"] = "invalid_auth" return False return True - else: - check_url = ( - f"{endpoint}/v1/models" if self._provider == API_PROVIDER_ANTHROPIC - else f"{endpoint}/models" - ) - async with session.get(check_url, headers=headers) as response: - if response.status == 401: - self._errors["base"] = "invalid_auth" - return False - elif response.status != 200: - self._errors["base"] = "cannot_connect" - return False - return True + session = async_get_clientsession(self.hass) + headers = build_auth_headers(self._provider, user_input[CONF_API_KEY]) + + from .providers import get_provider_config + check_path = get_provider_config(self._provider).get("check_path", "/models") + check_url = f"{endpoint}{check_path}" + + async with session.get(check_url, headers=headers) as response: + if response.status == 401: + self._errors["base"] = "invalid_auth" + return False + elif response.status != 200: + self._errors["base"] = "cannot_connect" + return False + return True except Exception as err: _LOGGER.error("API validation error: %s", str(err)) self._errors["base"] = "cannot_connect" return False - def _get_api_headers(self, user_input: Dict[str, Any]) -> Dict[str, str]: - """Get API headers based on provider.""" - if CONF_API_KEY not in user_input: - return {"Content-Type": "application/json"} - - api_key = user_input[CONF_API_KEY] - - if self._provider == API_PROVIDER_ANTHROPIC: - return { - "x-api-key": api_key, - "anthropic-version": "2023-06-01", - "Content-Type": "application/json" - } - elif self._provider == API_PROVIDER_GEMINI: - return { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json" - } - return { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json" - } - async def _create_entry(self, user_input: Dict[str, Any]) -> FlowResult: """Create the config entry with comprehensive data preservation.""" instance_name = user_input[CONF_NAME] @@ -436,10 +283,6 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): CONF_MAX_HISTORY_SIZE: user_input.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY), } - for key, value in user_input.items(): - if key not in entry_data: - entry_data[key] = value - _LOGGER.debug("Creating config entry with data: %s", safe_log_data(entry_data)) return self.async_create_entry( @@ -462,35 +305,13 @@ class OptionsFlowHandler(config_entries.OptionsFlow): self._errors = {} self._selected_provider = None - def _get_default_endpoint(self, provider: str) -> str: - """Get default endpoint for provider.""" - return get_default_endpoint(provider) - - def _get_default_model(self, provider: str) -> str: - """Get default model for provider.""" - return get_default_model(provider) - - def _get_api_headers(self, api_key: str, provider: str) -> Dict[str, str]: - """Get API headers based on provider.""" - if provider == API_PROVIDER_ANTHROPIC: - return { - "x-api-key": api_key, - "anthropic-version": "2023-06-01", - "Content-Type": "application/json" - } - return { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json" - } - async def _async_validate_api(self, provider: str, api_key: str, endpoint: str) -> bool: - """Validate API connection.""" + """Validate API connection using provider registry.""" try: if not api_key: self._errors["base"] = "invalid_auth" return False - # Validate endpoint URL (HTTPS-only, SSRF protection) try: endpoint = validate_endpoint(endpoint) except ValueError as err: @@ -498,17 +319,15 @@ class OptionsFlowHandler(config_entries.OptionsFlow): self._errors["base"] = "cannot_connect" return False - # For Gemini, just check if API key is present if provider == API_PROVIDER_GEMINI: return True session = async_get_clientsession(self.hass) - headers = self._get_api_headers(api_key, provider) + headers = build_auth_headers(provider, api_key) - check_url = ( - f"{endpoint}/v1/models" if provider == API_PROVIDER_ANTHROPIC - else f"{endpoint}/models" - ) + from .providers import get_provider_config + check_path = get_provider_config(provider).get("check_path", "/models") + check_url = f"{endpoint}{check_path}" async with session.get(check_url, headers=headers) as response: if response.status == 401: @@ -562,22 +381,45 @@ class OptionsFlowHandler(config_entries.OptionsFlow): # Use new defaults if provider changed, otherwise use current values if provider_changed: - default_endpoint = self._get_default_endpoint(provider) - default_model = self._get_default_model(provider) + default_endpoint = get_default_endpoint(provider) + default_model = get_default_model(provider) else: - default_endpoint = current_data.get(CONF_API_ENDPOINT, self._get_default_endpoint(provider)) - default_model = current_data.get(CONF_MODEL, self._get_default_model(provider)) + default_endpoint = current_data.get(CONF_API_ENDPOINT, get_default_endpoint(provider)) + default_model = current_data.get(CONF_MODEL, get_default_model(provider)) if user_input is not None: - # Validate API connection - api_key = user_input.get(CONF_API_KEY, current_data.get(CONF_API_KEY, "")) + api_key = user_input.get(CONF_API_KEY, "").strip() endpoint = user_input.get(CONF_API_ENDPOINT, default_endpoint) + # Require API key re-entry when endpoint or provider changed + stored_endpoint = current_data.get(CONF_API_ENDPOINT, "") + endpoint_changed = endpoint != stored_endpoint + if not api_key and (provider_changed or endpoint_changed): + self._errors["base"] = "invalid_auth" + return self.async_show_form( + step_id="settings", + data_schema=self._get_settings_schema( + provider=provider, + current_data=current_data, + user_input=user_input, + default_endpoint=default_endpoint, + default_model=default_model, + ), + errors=self._errors, + description_placeholders={ + "provider": provider + } + ) + + # Fall back to stored key if not re-entered and endpoint unchanged + if not api_key: + api_key = current_data.get(CONF_API_KEY, "") + if await self._async_validate_api(provider, api_key, endpoint): - # Merge with provider selection final_data = { CONF_API_PROVIDER: provider, - **user_input + **user_input, + CONF_API_KEY: api_key, } return self.async_create_entry(title="", data=final_data) @@ -591,7 +433,10 @@ class OptionsFlowHandler(config_entries.OptionsFlow): default_endpoint=default_endpoint, default_model=default_model, ), - errors=self._errors + errors=self._errors, + description_placeholders={ + "provider": provider + } ) return self.async_show_form( @@ -620,7 +465,7 @@ class OptionsFlowHandler(config_entries.OptionsFlow): data = user_input or current_data return vol.Schema({ - vol.Required(CONF_API_KEY): str, + vol.Optional(CONF_API_KEY, default=""): str, vol.Required( CONF_API_ENDPOINT, default=data.get(CONF_API_ENDPOINT, default_endpoint) diff --git a/custom_components/ha_text_ai/coordinator.py b/custom_components/ha_text_ai/coordinator.py index 998f819..6c8b977 100644 --- a/custom_components/ha_text_ai/coordinator.py +++ b/custom_components/ha_text_ai/coordinator.py @@ -8,58 +8,37 @@ The HA Text AI coordinator. """ from __future__ import annotations -import logging -import traceback -import aiofiles -import os -import json import asyncio -import math -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional, Union +import logging +import os +from datetime import timedelta +from typing import Any, Dict, List, Optional from homeassistant.core import HomeAssistant +from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from homeassistant.util import dt as dt_util -from homeassistant.exceptions import HomeAssistantError -from homeassistant.const import CONF_NAME -from .utils import normalize_name from .const import ( - DOMAIN, - STATE_READY, - STATE_PROCESSING, - STATE_ERROR, - STATE_RATE_LIMITED, - STATE_MAINTENANCE, - DEFAULT_MAX_TOKENS, - DEFAULT_TEMPERATURE, - DEFAULT_MAX_HISTORY, DEFAULT_API_TIMEOUT, DEFAULT_CONTEXT_MESSAGES, - ABSOLUTE_MAX_HISTORY_SIZE, - MAX_ATTRIBUTE_SIZE, - MAX_HISTORY_FILE_SIZE, + DEFAULT_MAX_HISTORY, + DEFAULT_MAX_TOKENS, + DEFAULT_TEMPERATURE, + STATE_ERROR, + STATE_MAINTENANCE, + STATE_PROCESSING, + STATE_RATE_LIMITED, + STATE_READY, TRUNCATION_INDICATOR, ) +from .history import HistoryManager +from .metrics import MetricsManager +from .utils import normalize_name _LOGGER = logging.getLogger(__name__) -class AsyncFileHandler: - """Async context manager for file operations.""" - - def __init__(self, file_path: str, mode: str = 'a'): - self.file_path = file_path - self.mode = mode - - async def __aenter__(self): - self.file = await aiofiles.open(self.file_path, self.mode) - return self.file - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.file.close() - class HATextAICoordinator(DataUpdateCoordinator): """Home Assistant Text AI Conversation Coordinator.""" @@ -79,53 +58,49 @@ class HATextAICoordinator(DataUpdateCoordinator): ) -> None: """Initialize coordinator.""" self.instance_name = instance_name - self.normalized_name = None - self.normalized_name = normalize_name(instance_name) - self._metrics_file = os.path.join( - hass.config.path(".storage"), - "ha_text_ai_history", - f"ha_text_ai_metrics_{normalize_name(instance_name)}.json" + + history_dir = os.path.join( + hass.config.path(".storage"), "ha_text_ai_history" + ) + metrics_file = os.path.join( + history_dir, + f"ha_text_ai_metrics_{self.normalized_name}.json", ) - # Initialize performance metrics first - self._performance_metrics = { - "total_tokens": 0, - "prompt_tokens": 0, - "completion_tokens": 0, - "successful_requests": 0, - "failed_requests": 0, - "total_errors": 0, - "average_latency": 0, - "max_latency": 0, - "min_latency": float("inf"), - } + # Delegate history and metrics to dedicated managers + self._history = HistoryManager( + hass=hass, + instance_name=instance_name, + normalized_name=self.normalized_name, + history_dir=history_dir, + max_history_size=max_history_size, + ) + self._metrics = MetricsManager( + hass=hass, + instance_name=instance_name, + metrics_file=metrics_file, + ) - # Continue with other initializations self.hass = hass self.client = client self.model = model self.temperature = temperature self.max_tokens = max_tokens - self.max_history_size = min( - max(1, max_history_size), - ABSOLUTE_MAX_HISTORY_SIZE - ) self.is_anthropic = is_anthropic self.api_timeout = api_timeout # Concurrency control self._request_lock = asyncio.Lock() - # Initialize essential attributes + # State flags self._is_processing = False self._is_rate_limited = False self._is_maintenance = False self.endpoint_status = "ready" - self._system_prompt = None - self._conversation_history = [] + self._system_prompt: Optional[str] = None - self._last_response = { + self._last_response: Dict[str, Any] = { "timestamp": dt_util.utcnow().isoformat(), "question": "", "response": "", @@ -135,605 +110,287 @@ class HATextAICoordinator(DataUpdateCoordinator): "error": None, } - update_interval_td = timedelta(seconds=update_interval) - - # Call super().__init__ super().__init__( hass, _LOGGER, name=instance_name, - update_interval=update_interval_td, + update_interval=timedelta(seconds=update_interval), ) - # Now initialize _initial_state - self._initial_state = { - "state": STATE_READY, - "metrics": self._performance_metrics.copy(), - "last_response": self.last_response.copy(), - "is_processing": self._is_processing, - "is_rate_limited": self._is_rate_limited, - "is_maintenance": self._is_maintenance, - "endpoint_status": self.endpoint_status, - "uptime": 0, - "system_prompt": self._system_prompt, - "history_size": len(self._conversation_history), - "conversation_history": self._conversation_history.copy(), - } - self.available = True self._state = STATE_READY - self._start_time = dt_util.utcnow() - - # Create history directory with safe mechanism - self._history_dir = os.path.join( - hass.config.path(".storage"), - "ha_text_ai_history" - ) - - # History file path using instance name — must be set before async_initialize - self._history_file = os.path.join( - self._history_dir, - f"{self.normalized_name}_history.json" - ) - - # Maximum history file size (1 MB) from const.py - self._max_history_file_size = MAX_HISTORY_FILE_SIZE - self.context_messages = context_messages - _LOGGER.info("Initialized HA Text AI coordinator with instance: %s", instance_name) + _LOGGER.info("Initialized HA Text AI coordinator: %s", instance_name) - async def async_initialize(self) -> None: - """Initialize coordinator: directories, history, metrics. Must be awaited.""" - await self._create_history_dir() - await self._check_history_directory() - await self._initialize_metrics() - await self.async_initialize_history_file() - await self._migrate_history_from_txt_to_json() + # ------------------------------------------------------------------ + # Convenience accessors for backward compatibility + # ------------------------------------------------------------------ + @property + def _conversation_history(self) -> List[Dict[str, Any]]: + return self._history.conversation_history + @property + def max_history_size(self) -> int: + return self._history.max_history_size + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + async def async_initialize(self) -> None: + """Initialize coordinator: directories, history, metrics. Must be awaited.""" + await self._history.async_initialize() + await self._metrics.async_initialize() + + async def async_shutdown(self) -> None: + """Shutdown coordinator.""" + _LOGGER.debug("Shutting down coordinator for %s", self.instance_name) + + # ------------------------------------------------------------------ + # Last response + # ------------------------------------------------------------------ @property def last_response(self) -> Dict[str, Any]: - """ - Get the last response with fallback to conversation history. - - Returns: - Dict containing last response information - """ - if self._last_response: - return self._last_response - - if self._conversation_history: - latest = self._conversation_history[-1] - return { - "timestamp": latest.get("timestamp"), - "question": latest.get("question", ""), - "response": latest.get("response", ""), - "model": self.model, - "instance": self.instance_name, - "normalized_name": self.normalized_name, - "error": None - } - - return { - "timestamp": dt_util.utcnow().isoformat(), - "question": "", - "response": "", - "model": self.model, - "instance": self.instance_name, - "normalized_name": self.normalized_name, - "error": None - } + """Get the last response.""" + return self._last_response @last_response.setter def last_response(self, value: Dict[str, Any]) -> None: - """ - Set the last response value. - - Args: - value: Dictionary containing response information - """ self._last_response = value - async def _file_exists(self, path: str) -> bool: - """ - Check if file exists asynchronously. - - Args: - path (str): Full path to the file to check - - Returns: - bool: True if file exists, False otherwise - """ + # ------------------------------------------------------------------ + # HA state update + # ------------------------------------------------------------------ + async def async_update_ha_state(self) -> None: + """Update Home Assistant state via coordinator refresh.""" try: - result = await self.hass.async_add_executor_job(os.path.exists, path) - - _LOGGER.debug(f"File existence check: {path} - {'Exists' if result else 'Not Found'}") - - return result - - except Exception as e: - _LOGGER.error(f"Error checking file existence for {path}: {e}") - return False - - async def _create_history_dir(self) -> None: - """ - Asynchronously create history directory. - """ - try: - def mkdir_sync(): - os.makedirs(self._history_dir, exist_ok=True) - - await self.hass.async_add_executor_job(mkdir_sync) - _LOGGER.debug(f"History directory created/verified: {self._history_dir}") - - except PermissionError: - _LOGGER.error(f"Permission denied when creating history directory: {self._history_dir}") - raise - except OSError as e: - _LOGGER.error(f"Error creating history directory {self._history_dir}: {e}") - raise - - async def _initialize_metrics(self) -> None: - self._performance_metrics = await self._load_metrics() or { - "total_tokens": 0, - "prompt_tokens": 0, - "completion_tokens": 0, - "successful_requests": 0, - "failed_requests": 0, - "total_errors": 0, - "average_latency": 0, - "max_latency": 0, - "min_latency": float("inf"), - } - - async def _load_metrics(self) -> Dict[str, Any]: - try: - if await self._file_exists(self._metrics_file): - def read_metrics(): - with open(self._metrics_file, 'r') as f: - try: - return json.load(f) - except json.JSONDecodeError: - _LOGGER.warning("Metrics file corrupted, creating new") - return None - return await self.hass.async_add_executor_job(read_metrics) - except Exception as e: - _LOGGER.warning(f"Failed to load metrics: {e}") - return None - - async def _save_metrics(self) -> None: - try: - def write_metrics(): - with open(self._metrics_file, 'w') as f: - json.dump(self._performance_metrics, f) - await self.hass.async_add_executor_job(write_metrics) - except Exception as e: - _LOGGER.warning(f"Failed to save metrics: {e}") - - async def _update_metrics(self, latency: float, response: dict) -> None: - """Update performance metrics and save to storage.""" - metrics = self._performance_metrics - tokens = response.get("tokens", {}) - - metrics["total_tokens"] += tokens.get("total", 0) - metrics["prompt_tokens"] += tokens.get("prompt", 0) - metrics["completion_tokens"] += tokens.get("completion", 0) - metrics["successful_requests"] += 1 - - metrics["average_latency"] = ( - (metrics["average_latency"] * (metrics["successful_requests"] - 1) + latency) - / metrics["successful_requests"] - ) - metrics["max_latency"] = max(metrics["max_latency"], latency) - metrics["min_latency"] = min(metrics["min_latency"], latency) - - # Save metrics after update - await self._save_metrics() - - async def _get_current_metrics(self) -> Dict[str, Any]: - """Get current performance metrics.""" - metrics = self._performance_metrics.copy() - _LOGGER.debug(f"Current performance metrics: {metrics}") - return metrics - - async def _handle_error(self, error: Exception) -> None: - """Enhanced error handling with metric persistence.""" - self._performance_metrics["total_errors"] += 1 - self._performance_metrics["failed_requests"] += 1 - - # Save metrics after error - await self._save_metrics() - - error_details = { - "timestamp": dt_util.utcnow().isoformat(), - "model": self.model, - "instance": self.instance_name, - "error_message": str(error), - "error_type": type(error).__name__, - "traceback": traceback.format_exc() if _LOGGER.isEnabledFor(logging.DEBUG) else None, - } - - # Specific error type handling - error_mapping = { - HomeAssistantError: {"is_ha_error": True}, - ConnectionError: { - "is_connection_error": True, - }, - TimeoutError: {"is_timeout": True}, - PermissionError: {"is_permission_denied": True}, - ValueError: {"is_validation_error": True} - } - - for error_type, error_flags in error_mapping.items(): - if isinstance(error, error_type): - error_details.update(error_flags) - break - - # Update system state based on error type - if error_details.get("is_rate_limited"): - self._is_rate_limited = True - _LOGGER.warning(f"Rate limit detected for {self.instance_name}") - - if error_details.get("is_connection_error"): - self.endpoint_status = "unavailable" - - self.last_response = error_details - _LOGGER.error(f"AI Processing Error: {error_details}") - - if _LOGGER.isEnabledFor(logging.DEBUG): - _LOGGER.debug(f"Full Error Traceback: {error_details['traceback']}") - - async def async_initialize_history_file(self) -> None: - """ - Asynchronously initialize history file and load existing history. - """ - try: - await self._create_history_dir() - - if await self._file_exists(self._history_file): - # Load existing history - async with AsyncFileHandler(self._history_file, 'r') as f: - content = await f.read() - if content: - history = json.loads(content) - # Validate and load history - if isinstance(history, list): - self._conversation_history = history[-self.max_history_size:] - _LOGGER.debug( - f"Loaded {len(self._conversation_history)} history entries " - f"for {self.instance_name}" - ) - else: - # Create new history file - async with AsyncFileHandler(self._history_file, 'w') as f: - await f.write(json.dumps([])) - - await self._check_history_size() - await self.async_update_ha_state() - - _LOGGER.debug(f"History file initialized: {self._history_file}") - - except Exception as e: - _LOGGER.error(f"Could not initialize history file: {e}") - _LOGGER.debug(traceback.format_exc()) - - # Size check to _update_history method - async def _update_history(self, question: str, response: dict) -> None: - """Update conversation history with size validation.""" - try: - history_entry = { - "timestamp": dt_util.utcnow().isoformat(), - "question": self._truncate_text(question, MAX_ATTRIBUTE_SIZE), - "response": self._truncate_text( - response.get("content", ""), - MAX_ATTRIBUTE_SIZE - ), - } - - # Check approximate entry size - entry_size = len(json.dumps(history_entry).encode('utf-8')) - current_size = await self._check_file_size(self._history_file) - - if current_size + entry_size > MAX_HISTORY_FILE_SIZE: - _LOGGER.warning( - f"History size limit approaching. " - f"Current: {current_size}, " - f"Entry: {entry_size}, " - f"Max: {MAX_HISTORY_FILE_SIZE}" - ) - await self._rotate_history() - - self._conversation_history.append(history_entry) - - # Enforce size limit - while len(self._conversation_history) > self.max_history_size: - self._conversation_history.pop(0) - - await self._write_history_entry(history_entry) - - except Exception as e: - _LOGGER.error(f"Error updating history: {e}") - _LOGGER.debug(traceback.format_exc()) - - async def _write_history_entry(self, entry: dict) -> None: - """Write history entry with file size checks.""" - try: - if not await self._file_exists(self._history_dir): - await self._create_history_dir() - - # Check current file size - current_size = 0 - if await self._file_exists(self._history_file): - current_size = await self.hass.async_add_executor_job( - os.path.getsize, - self._history_file - ) - - # Calculate approximate entry size - entry_size = len(json.dumps(entry).encode('utf-8')) - - if current_size + entry_size > MAX_HISTORY_FILE_SIZE: - _LOGGER.warning( - f"History file size limit reached. Current: {current_size}, " - f"Entry size: {entry_size}, Max: {MAX_HISTORY_FILE_SIZE}" - ) - # Trigger rotation before writing - await self._rotate_history() - - # Continue with writing - history = [] - if await self._file_exists(self._history_file): - async with AsyncFileHandler(self._history_file, 'r') as f: - content = await f.read() - if content: - history = json.loads(content) - - history.append(entry) - - if len(history) > self.max_history_size: - history = history[-self.max_history_size:] - - async with AsyncFileHandler(self._history_file, 'w') as f: - await f.write(json.dumps(history, indent=2)) - - except Exception as e: - _LOGGER.error(f"Error writing history entry: {e}") - _LOGGER.debug(traceback.format_exc()) - - async def _check_history_size(self) -> None: - """Verify and adjust history size if needed.""" - if len(self._conversation_history) > self.max_history_size: - _LOGGER.warning( - f"History size ({len(self._conversation_history)}) exceeds " - f"maximum ({self.max_history_size}). Trimming..." - ) - self._conversation_history = self._conversation_history[-self.max_history_size:] - - async def _check_file_size(self, file_path: str) -> int: - """ - Check file size asynchronously. - - Args: - file_path: Path to file to check - - Returns: - int: File size in bytes - """ - try: - if await self._file_exists(file_path): - size = await self.hass.async_add_executor_job(os.path.getsize, file_path) - _LOGGER.debug(f"File size check for {file_path}: {size} bytes") - return size - return 0 - except Exception as e: - _LOGGER.error(f"Error checking file size for {file_path}: {e}") - return 0 - - async def _rotate_history(self) -> None: - """Rotate conversation history with file management.""" - try: - _LOGGER.debug("Starting history rotation for %s", self._history_file) - await self._rotate_history_files() - _LOGGER.debug("Completed history rotation") - except Exception as e: - _LOGGER.error("Error rotating history: %s", e) - _LOGGER.debug(traceback.format_exc()) - - async def _check_history_directory(self) -> None: - """ - Asynchronously check history directory permissions and writability. - """ - try: - # First ensure directory exists - await self._create_history_dir() - - # Then test write permission in a separate thread - test_file_path = os.path.join(self._history_dir, ".write_test") - await self.hass.async_add_executor_job(self._sync_test_directory_write, test_file_path) - - except PermissionError: - _LOGGER.error(f"No write permissions for history directory: {self._history_dir}") - except Exception as e: - _LOGGER.error(f"Error checking history directory: {e}") - - def _sync_test_directory_write(self, test_file_path: str) -> None: - """ - Synchronous method to test directory write permissions. - """ - try: - with open(test_file_path, 'w') as f: - f.write("Permission test") - os.remove(test_file_path) - except Exception as e: - _LOGGER.error(f"Directory write test failed: {e}") - - async def _rotate_history_files(self) -> None: - """Rotate history files with size validation.""" - try: - if await self._file_exists(self._history_file): - current_size = await self._check_file_size(self._history_file) - - if current_size > MAX_HISTORY_FILE_SIZE: - _LOGGER.info( - f"Rotating history file. " - f"Current size: {current_size}, " - f"Max allowed: {MAX_HISTORY_FILE_SIZE}" - ) - - # Create archive filename with timestamp - archive_file = os.path.join( - self._history_dir, - f"{self.normalized_name}_history_{dt_util.utcnow().strftime('%Y%m%d_%H%M%S')}.json" - ) - - # Rotate files (shutil.move for cross-device safety) - import shutil - await self.hass.async_add_executor_job( - shutil.move, - self._history_file, - archive_file - ) - - # Create new history file with recent entries - async with AsyncFileHandler(self._history_file, 'w') as f: - await f.write(json.dumps( - self._conversation_history[-self.max_history_size:], - indent=2 - )) - - _LOGGER.info(f"History file rotated to: {archive_file}") - - except Exception as e: - _LOGGER.error(f"History rotation failed: {e}") - _LOGGER.debug(traceback.format_exc()) + await self.async_request_refresh() + except Exception as err: + _LOGGER.error("Error updating HA state for %s: %s", self.instance_name, err) async def _async_update_data(self) -> Dict[str, Any]: - """Update coordinator data with improved error handling and performance.""" + """Update coordinator data.""" try: current_state = self._get_current_state() - - # Get limited history with info - history_data = await self._get_limited_history() - - metrics = await self._get_current_metrics() - if metrics is None: - metrics = {} + history_data = self._history.get_limited_history() + metrics = await self._metrics.get_current_metrics() data = { "state": current_state, - "metrics": metrics, - "last_response": await self._get_sanitized_last_response(), + "metrics": metrics or {}, + "last_response": self._get_sanitized_last_response(), "is_processing": self._is_processing, "is_rate_limited": self._is_rate_limited, "is_maintenance": self._is_maintenance, "endpoint_status": self.endpoint_status, "uptime": self._calculate_uptime(), "system_prompt": self._get_truncated_system_prompt(), - "history_size": len(self._conversation_history), + "history_size": self._history.history_size, "conversation_history": history_data["entries"], "history_info": history_data["info"], "normalized_name": self.normalized_name, } - await self._validate_update_data(data) + self._validate_update_data(data) return data except Exception as err: _LOGGER.error("Error updating data: %s", err, exc_info=True) return self._get_safe_initial_state() - async def _get_limited_history(self) -> Dict[str, Any]: - """Get limited conversation history showing only last Q&A.""" - limited_history = [] + # ------------------------------------------------------------------ + # Question processing + # ------------------------------------------------------------------ + async def async_ask_question( + self, + question: str, + model: Optional[str] = None, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + system_prompt: Optional[str] = None, + context_messages: Optional[int] = None, + structured_output: bool = False, + json_schema: Optional[str] = None, + ) -> dict: + """Process a question with optional parameters.""" + return await self.async_process_question( + question, model, temperature, max_tokens, system_prompt, + context_messages, structured_output, json_schema, + ) - if self._conversation_history: - last_entry = self._conversation_history[-1] - limited_entry = { - "timestamp": last_entry["timestamp"], - "question": last_entry["question"][:4096] + (TRUNCATION_INDICATOR if len(last_entry["question"]) > 4096 else ""), - "response": last_entry["response"][:4096] + (TRUNCATION_INDICATOR if len(last_entry["response"]) > 4096 else ""), - } - limited_history.append(limited_entry) + async def async_process_question( + self, + question: str, + model: Optional[str] = None, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + system_prompt: Optional[str] = None, + context_messages: Optional[int] = None, + structured_output: bool = False, + json_schema: Optional[str] = None, + ) -> dict: + """Process question with context management.""" + if self.client is None: + raise HomeAssistantError("AI client not initialized") - history_info = { - "total_entries": len(self._conversation_history), - "displayed_entries": len(limited_history), - "full_history_available": True, - } + async with self._request_lock: + try: + self._is_processing = True + await self.async_update_ha_state() - return { - "entries": limited_history, - "full_history": self._conversation_history, - "info": history_info - } + temp_context = context_messages if context_messages is not None else self.context_messages + temp_model = model if model is not None else self.model + temp_temperature = temperature if temperature is not None else self.temperature + temp_max_tokens = max_tokens if max_tokens is not None else self.max_tokens + temp_system_prompt = system_prompt if system_prompt is not None else self._system_prompt - async def _get_sanitized_last_response(self) -> Dict[str, Any]: - """Get sanitized version of last response with truncation indicators.""" - response = self.last_response.copy() + start_time = dt_util.utcnow() - if "response" in response: - original_response = response["response"] - is_response_truncated = len(original_response) > 4096 - response["response"] = original_response[:4096] + (TRUNCATION_INDICATOR if is_response_truncated else "") - response["is_response_truncated"] = is_response_truncated - response["full_response_length"] = len(original_response) + messages = [] + if temp_system_prompt: + messages.append({"role": "system", "content": temp_system_prompt}) - if "question" in response: - original_question = response["question"] - is_question_truncated = len(original_question) > 4096 - response["question"] = original_question[:4096] + (TRUNCATION_INDICATOR if is_question_truncated else "") - response["is_question_truncated"] = is_question_truncated - response["full_question_length"] = len(original_question) + context_history = self._conversation_history[-temp_context:] + for entry in context_history: + messages.append({"role": "user", "content": entry["question"]}) + messages.append({"role": "assistant", "content": entry["response"]}) - return response + messages.append({"role": "user", "content": question}) - def _truncate_text(self, text: str, max_length: int = MAX_ATTRIBUTE_SIZE) -> str: - """Safely truncate text to maximum length.""" - return text[:max_length] if text else "" + response = await self._send_to_api( + question=question, + model=temp_model, + messages=messages, + temperature=temp_temperature, + max_tokens=temp_max_tokens, + structured_output=structured_output, + json_schema=json_schema, + ) - def _calculate_uptime(self) -> float: - """Calculate current uptime in seconds.""" - return (dt_util.utcnow() - self._start_time).total_seconds() + latency = (dt_util.utcnow() - start_time).total_seconds() + await self._metrics.update_metrics(latency, response) + await self._history.update_history(question, response) - def _get_truncated_system_prompt(self) -> Optional[str]: - """Get truncated system prompt.""" - return self._truncate_text(self._system_prompt, 4096) if self._system_prompt else None + return response - async def _validate_update_data(self, data: Dict[str, Any]) -> None: - """Validate update data structure and content.""" - required_keys = ["state", "metrics", "last_response"] - for key in required_keys: - if key not in data: - raise ValueError(f"Missing required key: {key}") + except Exception as err: + error_details = await self._metrics.handle_error(err, self.model) + if error_details.get("is_connection_error"): + self.endpoint_status = "unavailable" + self.last_response = error_details + raise HomeAssistantError(f"Failed to process question: {err}") - if not isinstance(data["metrics"], dict): - raise ValueError("Invalid metrics format") + finally: + self._is_processing = False + await self.async_update_ha_state() - async def async_update_ha_state(self) -> None: - """Update Home Assistant state via coordinator refresh.""" + async def _send_to_api( + self, + question: str, + model: str, + messages: List[Dict[str, str]], + temperature: float, + max_tokens: int, + structured_output: bool = False, + json_schema: Optional[str] = None, + ) -> dict: + """Send request to AI provider and return structured response.""" try: - _LOGGER.debug( - "Requesting state update for %s", self.instance_name, - ) - await self.async_request_refresh() - except Exception as err: - _LOGGER.error("Error updating HA state for %s: %s", self.instance_name, err) + async with asyncio.timeout(self.api_timeout): + response = await self.client.create( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + structured_output=structured_output, + json_schema=json_schema, + ) + # Reset error state on success + self._is_rate_limited = False + self.endpoint_status = "ready" + + timestamp = dt_util.utcnow().isoformat() + content = response["choices"][0]["message"]["content"] + tokens = { + "prompt": response["usage"]["prompt_tokens"], + "completion": response["usage"]["completion_tokens"], + "total": response["usage"]["total_tokens"], + } + + self.last_response = { + "timestamp": timestamp, + "question": question, + "response": content, + "model": model, + "instance": self.instance_name, + "normalized_name": self.normalized_name, + "error": None, + } + + return { + "content": content, + "tokens": tokens, + "model": model, + "timestamp": timestamp, + "instance": self.instance_name, + "question": question, + "success": True, + } + + except asyncio.TimeoutError: + raise HomeAssistantError("Request timed out") + except Exception as err: + _LOGGER.error("Error in API call: %s", err) + raise + + # ------------------------------------------------------------------ + # History / prompt delegation + # ------------------------------------------------------------------ + async def async_clear_history(self) -> None: + """Clear conversation history.""" + await self._history.async_clear_history() + await self.async_update_ha_state() + + async def async_get_history( + self, + limit: Optional[int] = None, + filter_model: Optional[str] = None, + start_date: Optional[str] = None, + include_metadata: bool = False, + sort_order: str = "newest", + ) -> List[Dict[str, Any]]: + """Get conversation history with optional filtering.""" + return await self._history.async_get_history( + limit=limit, + filter_model=filter_model, + start_date=start_date, + include_metadata=include_metadata, + sort_order=sort_order, + default_model=self.model, + ) + + async def async_set_system_prompt(self, prompt: str) -> None: + """Set system prompt.""" + self._system_prompt = prompt + await self.async_update_ha_state() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ def _get_current_state(self) -> str: - """Get current state based on internal flags.""" if self._is_processing: return STATE_PROCESSING - elif self._is_rate_limited: + if self._is_rate_limited: return STATE_RATE_LIMITED - elif self._is_maintenance: + if self._is_maintenance: return STATE_MAINTENANCE - elif self.last_response.get("error"): + if self.last_response.get("error") or self.last_response.get("error_message"): return STATE_ERROR return STATE_READY def _get_safe_initial_state(self) -> Dict[str, Any]: - """Return safe initial state when update fails.""" return { "state": STATE_ERROR, "metrics": {}, @@ -746,332 +403,44 @@ class HATextAICoordinator(DataUpdateCoordinator): "system_prompt": None, "history_size": 0, "conversation_history": [], + "history_info": { + "total_entries": 0, + "displayed_entries": 0, + "full_history_available": True, + }, "normalized_name": self.normalized_name, } - async def async_ask_question( - self, - question: str, - model: Optional[str] = None, - temperature: Optional[float] = None, - max_tokens: Optional[int] = None, - system_prompt: Optional[str] = None, - context_messages: Optional[int] = None, - structured_output: bool = False, - json_schema: Optional[str] = None, - ) -> dict: - """ - Process a question with optional parameters. + def _get_sanitized_last_response(self) -> Dict[str, Any]: + """Get sanitized version of last response with truncation.""" + response = self.last_response.copy() - This method is a direct wrapper around async_process_question, - allowing flexible AI interaction with optional model, temperature, - and context customization. - - Args: - question: The input question or prompt - model: Optional AI model to use - temperature: Optional response creativity level - max_tokens: Optional maximum response length - system_prompt: Optional system-level instruction - context_messages: Optional number of context messages to include - structured_output: Enable JSON structured output mode - json_schema: JSON Schema for structured output validation - - Returns: - Full response dictionary from the AI - """ - return await self.async_process_question( - question, model, temperature, max_tokens, system_prompt, context_messages, - structured_output, json_schema - ) - - async def async_process_question( - self, - question: str, - model: Optional[str] = None, - temperature: Optional[float] = None, - max_tokens: Optional[int] = None, - system_prompt: Optional[str] = None, - context_messages: Optional[int] = None, - structured_output: bool = False, - json_schema: Optional[str] = None, - ) -> dict: - """Process question with context management.""" - if self.client is None: - raise HomeAssistantError("AI client not initialized") - - async with self._request_lock: - try: - self._is_processing = True - await self.async_update_ha_state() - - temp_context_messages = context_messages if context_messages is not None else self.context_messages - temp_model = model if model is not None else self.model - temp_temperature = temperature if temperature is not None else self.temperature - temp_max_tokens = max_tokens if max_tokens is not None else self.max_tokens - temp_system_prompt = system_prompt if system_prompt is not None else self._system_prompt - - start_time = dt_util.utcnow() - - messages = [] - if temp_system_prompt: - messages.append({"role": "system", "content": temp_system_prompt}) - - context_history = self._conversation_history[-temp_context_messages:] - for entry in context_history: - messages.append({"role": "user", "content": entry["question"]}) - messages.append({"role": "assistant", "content": entry["response"]}) - - messages.append({"role": "user", "content": question}) - - kwargs = { - "model": temp_model, - "temperature": temp_temperature, - "max_tokens": temp_max_tokens, - "messages": messages, - "structured_output": structured_output, - "json_schema": json_schema, - } - - response = await self.async_process_message(question, **kwargs) - - end_time = dt_util.utcnow() - latency = (end_time - start_time).total_seconds() - await self._update_metrics(latency, response) - - await self._update_history(question, response) - - return response - - except Exception as err: - await self._handle_error(err) - raise HomeAssistantError(f"Failed to process question: {err}") - - finally: - self._is_processing = False - await self.async_update_ha_state() - - async def async_process_message(self, question: str, **kwargs) -> dict: - """Process message using the AI client.""" - try: - structured_output = kwargs.pop("structured_output", False) - json_schema = kwargs.pop("json_schema", None) - - async with asyncio.timeout(self.api_timeout): - # APIClient.create() handles provider routing internally - response = await self._process_openai_message( - question, structured_output=structured_output, - json_schema=json_schema, **kwargs + for field in ("response", "question"): + if field in response and response[field]: + original = response[field] + truncated = len(original) > 4096 + response[field] = ( + original[:4096] + TRUNCATION_INDICATOR if truncated else original ) + response[f"is_{field}_truncated"] = truncated + response[f"full_{field}_length"] = len(original) - # Reset error state on success - self._is_rate_limited = False - self.endpoint_status = "ready" + return response - # Add timestamp and model information to response - timestamp = dt_util.utcnow().isoformat() - model_used = kwargs.get("model", self.model) - - # Enhance response with metadata - enhanced_response = { - "content": response["content"], - "tokens": response.get("tokens", {}), - "model": model_used, - "timestamp": timestamp, - "instance": self.instance_name, - "question": question, - "success": True - } + def _calculate_uptime(self) -> float: + return (dt_util.utcnow() - self._start_time).total_seconds() - self.last_response = { - "timestamp": timestamp, - "question": question, - "response": response["content"], - "model": model_used, - "instance": self.instance_name, - "normalized_name": self.normalized_name, - "error": None, - } + def _get_truncated_system_prompt(self) -> Optional[str]: + if not self._system_prompt: + return None + if len(self._system_prompt) <= 4096: + return self._system_prompt + return self._system_prompt[:4096] + TRUNCATION_INDICATOR - return enhanced_response - - except asyncio.TimeoutError: - raise HomeAssistantError("Request timed out") - - except Exception as err: - await self._handle_error(err) - raise - - async def _process_openai_message(self, question: str, structured_output: bool = False, - json_schema: Optional[str] = None, **kwargs) -> dict: - """Process message using OpenAI API.""" - try: - response = await self.client.create( - model=kwargs["model"], - messages=kwargs["messages"], - temperature=kwargs["temperature"], - max_tokens=kwargs["max_tokens"], - structured_output=structured_output, - json_schema=json_schema, - ) - - return { - "content": response["choices"][0]["message"]["content"], - "tokens": { - "prompt": response["usage"]["prompt_tokens"], - "completion": response["usage"]["completion_tokens"], - "total": response["usage"]["total_tokens"], - }, - } - except Exception as e: - _LOGGER.error(f"Error in OpenAI API call: {str(e)}") - raise - - async def _migrate_history_from_txt_to_json(self) -> None: - try: - old_history_file = os.path.join( - self._history_dir, - f"{self.normalized_name}_history.txt" - ) - - if not await self._file_exists(old_history_file): - return - - _LOGGER.info(f"Found old history file for {self.instance_name}, starting migration to JSON format") - - # Read old history - history_entries = [] - async with AsyncFileHandler(old_history_file, 'r') as f: - content = await f.read() - - # Parse txt content - for line in content.split('\n'): - if not line or line.startswith("History initialized at:"): - continue - - try: - # Parse the old format: "timestamp: Question: question - Response: response" - parts = line.split(': ', 1) - if len(parts) != 2: - continue - - timestamp = parts[0] - content_parts = parts[1].split(' - ') - if len(content_parts) != 2: - continue - - question = content_parts[0].replace('Question: ', '') - response = content_parts[1].replace('Response: ', '') - - history_entries.append({ - "timestamp": timestamp, - "question": question, - "response": response - }) - - except Exception as e: - _LOGGER.warning(f"Error parsing history line: {line}. Error: {e}") - continue - - if history_entries: - # Write to new JSON file - async with AsyncFileHandler(self._history_file, 'w') as f: - await f.write(json.dumps(history_entries, indent=2)) - - # Backup old file (use shutil.move for cross-device safety) - import shutil - backup_file = old_history_file + '.backup' - await self.hass.async_add_executor_job(shutil.move, old_history_file, backup_file) - - _LOGGER.info( - f"Successfully migrated {len(history_entries)} entries " - f"from txt to JSON for {self.instance_name}. " - f"Old file backed up as {backup_file}" - ) - - # Update in-memory history - self._conversation_history = history_entries - - except Exception as e: - _LOGGER.error(f"Error during history migration for {self.instance_name}: {e}") - _LOGGER.debug(traceback.format_exc()) - - async def async_clear_history(self) -> None: - """ - Asynchronously clear conversation history. - - Removes in-memory history and deletes history file. - """ - try: - self._conversation_history = [] - if await self._file_exists(self._history_file): - await self.hass.async_add_executor_job(os.remove, self._history_file) - await self.async_update_ha_state() - _LOGGER.info(f"History for {self.instance_name} cleared successfully") - except Exception as e: - _LOGGER.error(f"Error clearing history: {e}") - _LOGGER.debug(traceback.format_exc()) - - async def async_get_history( - self, - limit: Optional[int] = None, - filter_model: Optional[str] = None, - start_date: Optional[str] = None, - include_metadata: bool = False, - sort_order: str = "newest" - ) -> List[Dict[str, Any]]: - """Get conversation history with optional filtering and sorting.""" - try: - history = self._conversation_history.copy() - - # Filter by model if specified - if filter_model: - history = [entry for entry in history if entry.get("model") == filter_model] - - # Filter by start date if specified - if start_date: - try: - from datetime import datetime - start_dt = datetime.fromisoformat(start_date.replace('Z', '+00:00')) - history = [ - entry for entry in history - if datetime.fromisoformat(entry["timestamp"].replace('Z', '+00:00')) >= start_dt - ] - except (ValueError, KeyError) as e: - _LOGGER.warning(f"Invalid start_date format: {start_date}. Error: {e}") - - # Sort history - if sort_order == "oldest": - history.sort(key=lambda x: x.get("timestamp", "")) - else: # newest (default) - history.sort(key=lambda x: x.get("timestamp", ""), reverse=True) - - # Apply limit - if limit and limit > 0: - history = history[:limit] - - # Add metadata if requested - if include_metadata: - for entry in history: - entry["metadata"] = { - "entry_size": len(str(entry)), - "question_length": len(entry.get("question", "")), - "response_length": len(entry.get("response", "")), - "model_used": entry.get("model", self.model), - "instance": self.instance_name - } - - return history - - except Exception as e: - _LOGGER.error(f"Error getting history: {e}") - return [] - - async def async_set_system_prompt(self, prompt: str) -> None: - """Set system prompt.""" - self._system_prompt = prompt - await self.async_update_ha_state() - - async def async_shutdown(self) -> None: - """Shutdown coordinator.""" - _LOGGER.debug("Shutting down coordinator for %s", self.instance_name) + @staticmethod + def _validate_update_data(data: Dict[str, Any]) -> None: + for key in ("state", "metrics", "last_response"): + if key not in data: + raise ValueError(f"Missing required key: {key}") + if not isinstance(data["metrics"], dict): + raise ValueError("Invalid metrics format") diff --git a/custom_components/ha_text_ai/history.py b/custom_components/ha_text_ai/history.py new file mode 100644 index 0000000..cec95dc --- /dev/null +++ b/custom_components/ha_text_ai/history.py @@ -0,0 +1,460 @@ +""" +History management for HA Text AI integration. + +@license: CC BY-NC-SA 4.0 International +@author: SMKRV +@github: https://github.com/smkrv/ha-text-ai +@source: https://github.com/smkrv/ha-text-ai +""" +from __future__ import annotations + +import json +import logging +import os +import shutil +import traceback +from typing import Any, Dict, List, Optional + +import aiofiles + +from homeassistant.core import HomeAssistant +from homeassistant.util import dt as dt_util + +from .const import ( + ABSOLUTE_MAX_HISTORY_SIZE, + MAX_ATTRIBUTE_SIZE, + MAX_HISTORY_FILE_SIZE, + TRUNCATION_INDICATOR, +) + +_LOGGER = logging.getLogger(__name__) + + +class AsyncFileHandler: + """Async context manager for file operations.""" + + def __init__(self, file_path: str, mode: str = "a"): + self.file_path = file_path + self.mode = mode + + async def __aenter__(self): + self.file = await aiofiles.open(self.file_path, self.mode) + return self.file + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.file.close() + + +class HistoryManager: + """Manages conversation history for an instance.""" + + def __init__( + self, + hass: HomeAssistant, + instance_name: str, + normalized_name: str, + history_dir: str, + max_history_size: int, + ) -> None: + self.hass = hass + self.instance_name = instance_name + self.normalized_name = normalized_name + self._history_dir = history_dir + self.max_history_size = min( + max(1, max_history_size), ABSOLUTE_MAX_HISTORY_SIZE + ) + self._history_file = os.path.join( + history_dir, f"{normalized_name}_history.json" + ) + self._max_history_file_size = MAX_HISTORY_FILE_SIZE + self._conversation_history: List[Dict[str, Any]] = [] + + @property + def conversation_history(self) -> List[Dict[str, Any]]: + return self._conversation_history + + @property + def history_size(self) -> int: + return len(self._conversation_history) + + async def async_initialize(self) -> None: + """Initialize history: directories, file, migration.""" + await self._create_history_dir() + await self._check_history_directory() + await self._initialize_history_file() + await self._migrate_history_from_txt_to_json() + + async def _file_exists(self, path: str) -> bool: + try: + return await self.hass.async_add_executor_job(os.path.exists, path) + except Exception as e: + _LOGGER.error("Error checking file existence for %s: %s", path, e) + return False + + async def _create_history_dir(self) -> None: + try: + await self.hass.async_add_executor_job( + os.makedirs, self._history_dir, 0o755, True + ) + except PermissionError: + _LOGGER.error("Permission denied creating history directory: %s", self._history_dir) + raise + except OSError as e: + _LOGGER.error("Error creating history directory %s: %s", self._history_dir, e) + raise + + async def _check_history_directory(self) -> None: + """Check history directory permissions and writability.""" + try: + await self._create_history_dir() + test_file_path = os.path.join(self._history_dir, ".write_test") + await self.hass.async_add_executor_job( + self._sync_test_directory_write, test_file_path + ) + except PermissionError: + _LOGGER.error("No write permissions for history directory: %s", self._history_dir) + except Exception as e: + _LOGGER.error("Error checking history directory: %s", e) + + @staticmethod + def _sync_test_directory_write(test_file_path: str) -> None: + try: + with open(test_file_path, "w") as f: + f.write("Permission test") + os.remove(test_file_path) + except Exception as e: + _LOGGER.error("Directory write test failed: %s", e) + + async def _initialize_history_file(self) -> None: + """Initialize history file and load existing history.""" + try: + await self._create_history_dir() + + if await self._file_exists(self._history_file): + async with AsyncFileHandler(self._history_file, "r") as f: + content = await f.read() + if content: + history = json.loads(content) + if isinstance(history, list): + self._conversation_history = history[ + -self.max_history_size : + ] + _LOGGER.debug( + "Loaded %d history entries for %s", + len(self._conversation_history), + self.instance_name, + ) + else: + async with AsyncFileHandler(self._history_file, "w") as f: + await f.write(json.dumps([])) + + await self._check_history_size() + except Exception as e: + _LOGGER.error("Could not initialize history file: %s", e) + _LOGGER.debug(traceback.format_exc()) + + async def update_history(self, question: str, response: dict) -> None: + """Update conversation history with size validation.""" + try: + history_entry = { + "timestamp": dt_util.utcnow().isoformat(), + "question": self._truncate_text(question, MAX_ATTRIBUTE_SIZE), + "response": self._truncate_text( + response.get("content", ""), MAX_ATTRIBUTE_SIZE + ), + } + + entry_size = len(json.dumps(history_entry).encode("utf-8")) + current_size = await self._check_file_size(self._history_file) + + if current_size + entry_size > MAX_HISTORY_FILE_SIZE: + _LOGGER.warning( + "History size limit approaching. Current: %d, Entry: %d, Max: %d", + current_size, entry_size, MAX_HISTORY_FILE_SIZE, + ) + await self._rotate_history() + + self._conversation_history.append(history_entry) + + while len(self._conversation_history) > self.max_history_size: + self._conversation_history.pop(0) + + await self._write_history_entry(history_entry) + except Exception as e: + _LOGGER.error("Error updating history: %s", e) + _LOGGER.debug(traceback.format_exc()) + + async def _write_history_entry(self, entry: dict) -> None: + """Write history entry with file size checks.""" + try: + if not await self._file_exists(self._history_dir): + await self._create_history_dir() + + current_size = 0 + if await self._file_exists(self._history_file): + current_size = await self.hass.async_add_executor_job( + os.path.getsize, self._history_file + ) + + entry_size = len(json.dumps(entry).encode("utf-8")) + if current_size + entry_size > MAX_HISTORY_FILE_SIZE: + _LOGGER.warning( + "History file size limit reached. Current: %d, Entry: %d, Max: %d", + current_size, entry_size, MAX_HISTORY_FILE_SIZE, + ) + await self._rotate_history() + + history = [] + if await self._file_exists(self._history_file): + async with AsyncFileHandler(self._history_file, "r") as f: + content = await f.read() + if content: + history = json.loads(content) + + history.append(entry) + if len(history) > self.max_history_size: + history = history[-self.max_history_size :] + + async with AsyncFileHandler(self._history_file, "w") as f: + await f.write(json.dumps(history, indent=2)) + except Exception as e: + _LOGGER.error("Error writing history entry: %s", e) + _LOGGER.debug(traceback.format_exc()) + + async def _check_history_size(self) -> None: + if len(self._conversation_history) > self.max_history_size: + _LOGGER.warning( + "History size (%d) exceeds maximum (%d). Trimming...", + len(self._conversation_history), self.max_history_size, + ) + self._conversation_history = self._conversation_history[ + -self.max_history_size : + ] + + async def _check_file_size(self, file_path: str) -> int: + try: + if await self._file_exists(file_path): + return await self.hass.async_add_executor_job( + os.path.getsize, file_path + ) + return 0 + except Exception as e: + _LOGGER.error("Error checking file size for %s: %s", file_path, e) + return 0 + + async def _rotate_history(self) -> None: + try: + _LOGGER.debug("Starting history rotation for %s", self._history_file) + await self._rotate_history_files() + except Exception as e: + _LOGGER.error("Error rotating history: %s", e) + _LOGGER.debug(traceback.format_exc()) + + async def _rotate_history_files(self) -> None: + """Rotate history files with size validation.""" + try: + if await self._file_exists(self._history_file): + current_size = await self._check_file_size(self._history_file) + + if current_size > MAX_HISTORY_FILE_SIZE: + _LOGGER.info( + "Rotating history file. Current size: %d, Max: %d", + current_size, MAX_HISTORY_FILE_SIZE, + ) + + archive_file = os.path.join( + self._history_dir, + f"{self.normalized_name}_history_{dt_util.utcnow().strftime('%Y%m%d_%H%M%S')}.json", + ) + + await self.hass.async_add_executor_job( + shutil.move, self._history_file, archive_file + ) + + async with AsyncFileHandler(self._history_file, "w") as f: + await f.write( + json.dumps( + self._conversation_history[ + -self.max_history_size : + ], + indent=2, + ) + ) + + _LOGGER.info("History file rotated to: %s", archive_file) + except Exception as e: + _LOGGER.error("History rotation failed: %s", e) + _LOGGER.debug(traceback.format_exc()) + + async def _migrate_history_from_txt_to_json(self) -> None: + """Migrate old .txt history to .json format.""" + try: + old_history_file = os.path.join( + self._history_dir, f"{self.normalized_name}_history.txt" + ) + + if not await self._file_exists(old_history_file): + return + + # Skip migration if JSON history already has entries + if self._conversation_history: + _LOGGER.debug( + "JSON history already has %d entries for %s, skipping txt migration", + len(self._conversation_history), self.instance_name, + ) + return + + _LOGGER.info( + "Found old history file for %s, migrating to JSON", self.instance_name + ) + + history_entries = [] + async with AsyncFileHandler(old_history_file, "r") as f: + content = await f.read() + + for line in content.split("\n"): + if not line or line.startswith("History initialized at:"): + continue + try: + parts = line.split(": ", 1) + if len(parts) != 2: + continue + timestamp = parts[0] + content_parts = parts[1].split(" - ") + if len(content_parts) != 2: + continue + question = content_parts[0].replace("Question: ", "") + response = content_parts[1].replace("Response: ", "") + history_entries.append( + { + "timestamp": timestamp, + "question": question, + "response": response, + } + ) + except Exception as e: + _LOGGER.warning("Error parsing history line: %s. Error: %s", line, e) + continue + + if history_entries: + async with AsyncFileHandler(self._history_file, "w") as f: + await f.write(json.dumps(history_entries, indent=2)) + + backup_file = old_history_file + ".backup" + await self.hass.async_add_executor_job( + shutil.move, old_history_file, backup_file + ) + + _LOGGER.info( + "Migrated %d entries from txt to JSON for %s. Old file: %s", + len(history_entries), self.instance_name, backup_file, + ) + + self._conversation_history = history_entries + except Exception as e: + _LOGGER.error("Error during history migration for %s: %s", self.instance_name, e) + _LOGGER.debug(traceback.format_exc()) + + async def async_clear_history(self) -> None: + """Clear conversation history.""" + try: + self._conversation_history = [] + if await self._file_exists(self._history_file): + await self.hass.async_add_executor_job(os.remove, self._history_file) + _LOGGER.info("History for %s cleared", self.instance_name) + except Exception as e: + _LOGGER.error("Error clearing history: %s", e) + _LOGGER.debug(traceback.format_exc()) + + async def async_get_history( + self, + limit: Optional[int] = None, + filter_model: Optional[str] = None, + start_date: Optional[str] = None, + include_metadata: bool = False, + sort_order: str = "newest", + default_model: str = "", + ) -> List[Dict[str, Any]]: + """Get conversation history with optional filtering and sorting.""" + try: + history = self._conversation_history.copy() + + if filter_model: + history = [ + entry for entry in history if entry.get("model") == filter_model + ] + + if start_date: + try: + from datetime import datetime + + start_dt = datetime.fromisoformat( + start_date.replace("Z", "+00:00") + ) + history = [ + entry + for entry in history + if datetime.fromisoformat( + entry["timestamp"].replace("Z", "+00:00") + ) + >= start_dt + ] + except (ValueError, KeyError) as e: + _LOGGER.warning("Invalid start_date format: %s. Error: %s", start_date, e) + + if sort_order == "oldest": + history.sort(key=lambda x: x.get("timestamp", "")) + else: + history.sort(key=lambda x: x.get("timestamp", ""), reverse=True) + + if limit and limit > 0: + history = history[:limit] + + if include_metadata: + for entry in history: + entry["metadata"] = { + "entry_size": len(str(entry)), + "question_length": len(entry.get("question", "")), + "response_length": len(entry.get("response", "")), + "model_used": entry.get("model", default_model), + "instance": self.instance_name, + } + + return history + except Exception as e: + _LOGGER.error("Error getting history: %s", e) + return [] + + def get_limited_history(self) -> Dict[str, Any]: + """Get limited conversation history showing only last Q&A.""" + limited_history = [] + + if self._conversation_history: + last_entry = self._conversation_history[-1] + limited_entry = { + "timestamp": last_entry["timestamp"], + "question": self._truncate_text(last_entry["question"], 4096), + "response": self._truncate_text(last_entry["response"], 4096), + } + limited_history.append(limited_entry) + + history_info = { + "total_entries": len(self._conversation_history), + "displayed_entries": len(limited_history), + "full_history_available": True, + } + + return { + "entries": limited_history, + "full_history": list(self._conversation_history), + "info": history_info, + } + + @staticmethod + def _truncate_text(text: str, max_length: int = MAX_ATTRIBUTE_SIZE) -> str: + """Safely truncate text to maximum length with indicator.""" + if not text: + return "" + if len(text) <= max_length: + return text + return text[:max_length] + TRUNCATION_INDICATOR diff --git a/custom_components/ha_text_ai/metrics.py b/custom_components/ha_text_ai/metrics.py new file mode 100644 index 0000000..31f7503 --- /dev/null +++ b/custom_components/ha_text_ai/metrics.py @@ -0,0 +1,152 @@ +""" +Metrics management for HA Text AI integration. + +@license: CC BY-NC-SA 4.0 International +@author: SMKRV +@github: https://github.com/smkrv/ha-text-ai +@source: https://github.com/smkrv/ha-text-ai +""" +from __future__ import annotations + +import json +import logging +import os +import traceback +from typing import Any, Dict + +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import HomeAssistantError +from homeassistant.util import dt as dt_util + +_LOGGER = logging.getLogger(__name__) + +DEFAULT_METRICS: Dict[str, Any] = { + "total_tokens": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "successful_requests": 0, + "failed_requests": 0, + "total_errors": 0, + "average_latency": 0, + "max_latency": 0, + "min_latency": 0, +} + + +class MetricsManager: + """Manages performance metrics for an instance.""" + + def __init__( + self, + hass: HomeAssistant, + instance_name: str, + metrics_file: str, + ) -> None: + self.hass = hass + self.instance_name = instance_name + self._metrics_file = metrics_file + self._performance_metrics: Dict[str, Any] = DEFAULT_METRICS.copy() + + @property + def metrics(self) -> Dict[str, Any]: + return self._performance_metrics + + async def async_initialize(self) -> None: + """Load metrics from storage or create defaults.""" + loaded = await self._load_metrics() + self._performance_metrics = loaded or DEFAULT_METRICS.copy() + + async def _load_metrics(self) -> Dict[str, Any] | None: + try: + exists = await self.hass.async_add_executor_job( + os.path.exists, self._metrics_file + ) + if exists: + def read_metrics(): + with open(self._metrics_file, "r") as f: + try: + return json.load(f) + except json.JSONDecodeError: + _LOGGER.warning("Metrics file corrupted, creating new") + return None + + return await self.hass.async_add_executor_job(read_metrics) + except Exception as e: + _LOGGER.warning("Failed to load metrics: %s", e) + return None + + async def _save_metrics(self) -> None: + try: + def write_metrics(): + with open(self._metrics_file, "w") as f: + json.dump(self._performance_metrics, f) + + await self.hass.async_add_executor_job(write_metrics) + except Exception as e: + _LOGGER.warning("Failed to save metrics: %s", e) + + async def update_metrics(self, latency: float, response: dict) -> None: + """Update performance metrics after a successful request.""" + metrics = self._performance_metrics + tokens = response.get("tokens", {}) + + metrics["total_tokens"] += tokens.get("total", 0) + metrics["prompt_tokens"] += tokens.get("prompt", 0) + metrics["completion_tokens"] += tokens.get("completion", 0) + metrics["successful_requests"] += 1 + + metrics["average_latency"] = ( + (metrics["average_latency"] * (metrics["successful_requests"] - 1) + latency) + / metrics["successful_requests"] + ) + metrics["max_latency"] = max(metrics["max_latency"], latency) + if metrics["min_latency"] == 0: + metrics["min_latency"] = latency + else: + metrics["min_latency"] = min(metrics["min_latency"], latency) + + await self._save_metrics() + + async def get_current_metrics(self) -> Dict[str, Any]: + """Get current performance metrics.""" + return self._performance_metrics.copy() + + async def handle_error( + self, + error: Exception, + model: str, + ) -> Dict[str, Any]: + """Record an error in metrics and return error details.""" + self._performance_metrics["total_errors"] += 1 + self._performance_metrics["failed_requests"] += 1 + await self._save_metrics() + + error_details: Dict[str, Any] = { + "timestamp": dt_util.utcnow().isoformat(), + "model": model, + "instance": self.instance_name, + "error_message": str(error), + "error_type": type(error).__name__, + "traceback": traceback.format_exc() + if _LOGGER.isEnabledFor(logging.DEBUG) + else None, + } + + error_mapping = { + HomeAssistantError: {"is_ha_error": True}, + ConnectionError: {"is_connection_error": True}, + TimeoutError: {"is_timeout": True}, + PermissionError: {"is_permission_denied": True}, + ValueError: {"is_validation_error": True}, + } + + for error_type, error_flags in error_mapping.items(): + if isinstance(error, error_type): + error_details.update(error_flags) + break + + _LOGGER.error("AI Processing Error: %s", error_details) + if _LOGGER.isEnabledFor(logging.DEBUG): + _LOGGER.debug("Full Error Traceback: %s", error_details.get("traceback")) + + return error_details diff --git a/custom_components/ha_text_ai/sensor.py b/custom_components/ha_text_ai/sensor.py index 31b4d67..1222b89 100644 --- a/custom_components/ha_text_ai/sensor.py +++ b/custom_components/ha_text_ai/sensor.py @@ -9,8 +9,6 @@ Sensor platform for HA Text AI. import logging import math from typing import Any, Dict -from datetime import datetime, timedelta - from homeassistant.components.sensor import ( SensorEntity, SensorEntityDescription, @@ -277,9 +275,7 @@ class HATextAISensor(CoordinatorEntity, SensorEntity): METRIC_FAILED_REQUESTS: metrics.get("failed_requests", 0), METRIC_AVERAGE_LATENCY: round(metrics.get("average_latency", 0), 2), METRIC_MAX_LATENCY: round(metrics.get("max_latency", 0), 2), - METRIC_MIN_LATENCY: (metrics.get("min_latency") - if metrics.get("min_latency") != float("inf") - else None), + METRIC_MIN_LATENCY: metrics.get("min_latency", 0) or None, }) # Last response handling diff --git a/custom_components/ha_text_ai/utils.py b/custom_components/ha_text_ai/utils.py index 66331b3..1b126ea 100644 --- a/custom_components/ha_text_ai/utils.py +++ b/custom_components/ha_text_ai/utils.py @@ -7,6 +7,7 @@ Utility functions for HA Text AI integration. @source: https://github.com/smkrv/ha-text-ai """ import ipaddress +import socket from typing import Any from urllib.parse import urlparse @@ -47,7 +48,7 @@ def validate_endpoint(endpoint: str) -> str: if not hostname: raise ValueError("Invalid endpoint URL: no hostname") - # Block private/reserved IPs + # Block private/reserved IPs (direct IP or resolved hostname) try: addr = ipaddress.ip_address(hostname) if addr.is_private or addr.is_reserved or addr.is_loopback or addr.is_link_local: @@ -55,6 +56,23 @@ def validate_endpoint(endpoint: str) -> str: except ValueError as e: if "not allowed" in str(e): raise - # Not an IP address — it's a hostname, which is fine + # Not an IP literal — resolve hostname and check all resolved IPs + # to prevent DNS rebinding attacks + try: + addrinfos = socket.getaddrinfo(hostname, None) + for family, _type, _proto, _canonname, sockaddr in addrinfos: + ip_str = sockaddr[0] + resolved_addr = ipaddress.ip_address(ip_str) + if ( + resolved_addr.is_private + or resolved_addr.is_reserved + or resolved_addr.is_loopback + or resolved_addr.is_link_local + ): + raise ValueError( + f"Hostname {hostname} resolves to private/reserved IP {ip_str}" + ) + except socket.gaierror: + raise ValueError(f"Cannot resolve hostname: {hostname}") return endpoint.rstrip("/")